{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|default_exp models.layers"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Layers"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ">Helper functions used to build PyTorch timeseries models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "from torch.jit import TracerWarning\n",
    "import warnings\n",
    "from torch.nn.utils import weight_norm, spectral_norm\n",
    "from torch.nn.init import normal_\n",
    "from fastcore.basics import snake2camel\n",
    "from fastcore.test import test_eq\n",
    "from fastai.layers import *\n",
    "from fastai.losses import *\n",
    "from tsai.imports import *\n",
    "from tsai.utils import *\n",
    "\n",
    "warnings.filterwarnings(\"ignore\", category=TracerWarning)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def test_module_to_torchscript(\n",
    "    m:torch.nn.Module, # The PyTorch module to be tested.\n",
    "    inputs:Tensor, # A tensor or tuple of tensors representing the inputs to the model.\n",
    "    trace:bool=True, #  If `True`, attempts to trace the model. Defaults to `True`.\n",
    "    script:bool=True, # If `True`, attempts to script the model. Defaults to `True`.\n",
    "    serialize:bool=True, # If `True`, saves and loads the traced/scripted module to ensure it can be serialized. Defaults to `True`.\n",
    "    verbose:bool=True, # If `True`, prints detailed information about the tracing and scripting process. Defaults to `True`.\n",
    "):\n",
    "    \"Tests if a PyTorch module can be correctly traced or scripted and serialized\"\n",
    "\n",
    "    m = m.eval()\n",
    "    m_name = m.__class__.__name__\n",
    "\n",
    "    # Ensure inputs are in a tuple or list format\n",
    "    inp_is_tuple = isinstance(inputs, (tuple, list))\n",
    "\n",
    "    # Get the model's output\n",
    "    output = m(*inputs) if inp_is_tuple else m(inputs)\n",
    "    output_shapes = output.shape if not isinstance(output, (tuple, list)) else [o.shape for o in output]\n",
    "    if verbose:\n",
    "        print(f\"output.shape: {output_shapes}\")\n",
    "\n",
    "    # Try tracing the model\n",
    "    if trace:\n",
    "        if verbose:\n",
    "            print(\"Tracing...\")\n",
    "        try:\n",
    "            traced_m = torch.jit.trace(m, inputs)\n",
    "            if serialize:\n",
    "                file_path = Path(f\"test_traced_{m_name}.pt\")\n",
    "                torch.jit.save(traced_m, file_path)\n",
    "                traced_mod = torch.jit.load(file_path)\n",
    "                file_path.unlink()\n",
    "            traced_output = traced_m(*inputs) if inp_is_tuple else traced_m(inputs)\n",
    "            torch.testing.assert_close(traced_output, output)\n",
    "            if verbose:\n",
    "                print(f\"...{m_name} has been successfully traced 😃\\n\")\n",
    "            return True\n",
    "        except Exception as e:\n",
    "            if verbose:\n",
    "                print(f\"{m_name} cannot be traced 😔\")\n",
    "                print(e)\n",
    "                print(\"\\n\")\n",
    "\n",
    "    # Try scripting the model\n",
    "    if script:\n",
    "        if verbose:\n",
    "            print(\"Scripting...\")\n",
    "        try:\n",
    "            scripted_m = torch.jit.script(m)\n",
    "            if serialize:\n",
    "                file_path = Path(f\"test_scripted_{m_name}.pt\")\n",
    "                torch.jit.save(scripted_m, file_path)\n",
    "                scripted_mod = torch.jit.load(file_path)\n",
    "                file_path.unlink()\n",
    "            scripted_output = scripted_m(*inputs) if inp_is_tuple else scripted_m(inputs)\n",
    "            torch.testing.assert_close(scripted_output, output)\n",
    "            if verbose:\n",
    "                print(f\"...{m_name} has been successfully scripted 😃\\n\")\n",
    "            return True\n",
    "        except Exception as e:\n",
    "            if verbose:\n",
    "                print(f\"{m_name} cannot be scripted 😔\")\n",
    "                print(e)\n",
    "\n",
    "    return False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "output.shape: torch.Size([3, 2])\n",
      "Tracing...\n",
      "...Linear has been successfully traced 😃\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m = nn.Linear(10, 2)\n",
    "inp = torch.randn(3, 10)\n",
    "test_module_to_torchscript(m, inp, trace=True, script=True, serialize=True, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def init_lin_zero(m):\n",
    "    if isinstance(m, (nn.Linear)):\n",
    "        if getattr(m, 'bias', None) is not None: nn.init.constant_(m.bias, 0)\n",
    "        nn.init.constant_(m.weight, 0)\n",
    "    for l in m.children(): init_lin_zero(l)\n",
    "\n",
    "lin_zero_init = init_lin_zero"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class SwishBeta(Module):\n",
    "    def __multiinit__(self, beta=1.):\n",
    "        self.sigmoid = torch.sigmoid\n",
    "        self.beta = nn.Parameter(torch.Tensor(1).fill_(beta))\n",
    "    def forward(self, x): return x.mul(self.sigmoid(x*self.beta))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class SmeLU(nn.Module):\n",
    "    \"Smooth ReLU activation function based on https://arxiv.org/pdf/2202.06499.pdf\"\n",
    "\n",
    "    def __init__(self,\n",
    "        beta: float = 2. # Beta value\n",
    "        ) -> None:\n",
    "        super().__init__()\n",
    "        self.beta = abs(beta)\n",
    "\n",
    "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
    "        return torch.where(torch.abs(x) <= self.beta, ((x + self.beta) ** 2) / (4. * self.beta), F.relu(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class Chomp1d(nn.Module):\n",
    "    def __init__(self, chomp_size):\n",
    "        super(Chomp1d, self).__init__()\n",
    "        self.chomp_size = chomp_size\n",
    "\n",
    "    def forward(self, x):\n",
    "        return x[:, :, :-self.chomp_size].contiguous()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def same_padding1d(seq_len, ks, stride=1, dilation=1):\n",
    "    \"Same padding formula as used in Tensorflow\"\n",
    "    p = (seq_len - 1) * stride + (ks - 1) * dilation + 1 - seq_len\n",
    "    return p // 2, p - p // 2\n",
    "\n",
    "\n",
    "class Pad1d(nn.ConstantPad1d):\n",
    "    def __init__(self, padding, value=0.):\n",
    "        super().__init__(padding, value)\n",
    "\n",
    "\n",
    "# @delegates(nn.Conv1d.__init__)\n",
    "class SameConv1d(Module):\n",
    "    \"Conv1d with padding='same'\"\n",
    "    def __init__(self, ni, nf, ks=3, stride=1, dilation=1, **kwargs):\n",
    "        self.ks, self.stride, self.dilation = ks, stride, dilation\n",
    "        self.conv1d_same = nn.Conv1d(ni, nf, ks, stride=stride, dilation=dilation, **kwargs)\n",
    "        self.weight = self.conv1d_same.weight\n",
    "        self.bias = self.conv1d_same.bias\n",
    "        self.pad = Pad1d\n",
    "\n",
    "    def forward(self, x):\n",
    "        self.padding = same_padding1d(x.shape[-1], self.ks, dilation=self.dilation) #stride=self.stride not used in padding calculation!\n",
    "        return self.conv1d_same(self.pad(self.padding)(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def same_padding2d(H, W, ks, stride=(1, 1), dilation=(1, 1)):\n",
    "    \"Same padding formula as used in Tensorflow\"\n",
    "    if isinstance(ks, Integral): ks = (ks, ks)\n",
    "    if ks[0] == 1:  p_h = 0\n",
    "    else:  p_h = (H - 1) * stride[0] + (ks[0] - 1) * dilation[0] + 1 - H\n",
    "    if ks[1] == 1:  p_w = 0\n",
    "    else:  p_w = (W - 1) * stride[1] + (ks[1] - 1) * dilation[1] + 1 - W\n",
    "    return (p_w // 2, p_w - p_w // 2, p_h // 2, p_h - p_h // 2)\n",
    "\n",
    "\n",
    "class Pad2d(nn.ConstantPad2d):\n",
    "    def __init__(self, padding, value=0.):\n",
    "        super().__init__(padding, value)\n",
    "\n",
    "\n",
    "# @delegates(nn.Conv2d.__init__)\n",
    "class Conv2dSame(Module):\n",
    "    \"Conv2d with padding='same'\"\n",
    "    def __init__(self, ni, nf, ks=(3, 3), stride=(1, 1), dilation=(1, 1), **kwargs):\n",
    "        if isinstance(ks, Integral): ks = (ks, ks)\n",
    "        if isinstance(stride, Integral): stride = (stride, stride)\n",
    "        if isinstance(dilation, Integral): dilation = (dilation, dilation)\n",
    "        self.ks, self.stride, self.dilation = ks, stride, dilation\n",
    "        self.conv2d_same = nn.Conv2d(ni, nf, ks, stride=stride, dilation=dilation, **kwargs)\n",
    "        self.weight = self.conv2d_same.weight\n",
    "        self.bias = self.conv2d_same.bias\n",
    "        self.pad = Pad2d\n",
    "\n",
    "    def forward(self, x):\n",
    "        self.padding = same_padding2d(x.shape[-2], x.shape[-1], self.ks, dilation=self.dilation) #stride=self.stride not used in padding calculation!\n",
    "        return self.conv2d_same(self.pad(self.padding)(x))\n",
    "\n",
    "\n",
    "# @delegates(nn.Conv2d.__init__)\n",
    "def Conv2d(ni, nf, kernel_size=None, ks=None, stride=1, padding='same', dilation=1, init='auto', bias_std=0.01, **kwargs):\n",
    "    \"conv1d layer with padding='same', 'valid', or any integer (defaults to 'same')\"\n",
    "    assert not (kernel_size and ks), 'use kernel_size or ks but not both simultaneously'\n",
    "    assert kernel_size is not None or ks is not None, 'you need to pass a ks'\n",
    "    kernel_size = kernel_size or ks\n",
    "    if padding == 'same':\n",
    "        conv = Conv2dSame(ni, nf, kernel_size, stride=stride, dilation=dilation, **kwargs)\n",
    "    elif padding == 'valid': conv = nn.Conv2d(ni, nf, kernel_size, stride=stride, padding=0, dilation=dilation, **kwargs)\n",
    "    else: conv = nn.Conv2d(ni, nf, kernel_size, stride=stride, padding=padding, dilation=dilation, **kwargs)\n",
    "    init_linear(conv, None, init=init, bias_std=bias_std)\n",
    "    return conv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bs = 2\n",
    "c_in = 3\n",
    "c_out = 5\n",
    "h = 16\n",
    "w = 20\n",
    "t = torch.rand(bs, c_in, h, w)\n",
    "test_eq(Conv2dSame(c_in, c_out, ks=3, stride=1, dilation=1, bias=False)(t).shape, (bs, c_out, h, w))\n",
    "test_eq(Conv2dSame(c_in, c_out, ks=(3, 1), stride=1, dilation=1, bias=False)(t).shape, (bs, c_out, h, w))\n",
    "test_eq(Conv2dSame(c_in, c_out, ks=3, stride=(1, 1), dilation=(2, 2), bias=False)(t).shape, (bs, c_out, h, w))\n",
    "test_eq(Conv2dSame(c_in, c_out, ks=3, stride=(2, 2), dilation=(1, 1), bias=False)(t).shape, (bs, c_out, h//2, w//2))\n",
    "test_eq(Conv2dSame(c_in, c_out, ks=3, stride=(2, 2), dilation=(2, 2), bias=False)(t).shape, (bs, c_out, h//2, w//2))\n",
    "test_eq(Conv2d(c_in, c_out, ks=3, padding='same', stride=1, dilation=1, bias=False)(t).shape, (bs, c_out, h, w))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class CausalConv1d(torch.nn.Conv1d):\n",
    "    def __init__(self, ni, nf, ks, stride=1, dilation=1, groups=1, bias=True):\n",
    "        super(CausalConv1d, self).__init__(ni, nf, kernel_size=ks, stride=stride, padding=0, dilation=dilation, groups=groups, bias=bias)\n",
    "        self.__padding = (ks - 1) * dilation\n",
    "    def forward(self, input):\n",
    "        return super(CausalConv1d, self).forward(F.pad(input, (self.__padding, 0)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "# @delegates(nn.Conv1d.__init__)\n",
    "def Conv1d(ni, nf, kernel_size=None, ks=None, stride=1, padding='same', dilation=1, init='auto', bias_std=0.01, **kwargs):\n",
    "    \"conv1d layer with padding='same', 'causal', 'valid', or any integer (defaults to 'same')\"\n",
    "    assert not (kernel_size and ks), 'use kernel_size or ks but not both simultaneously'\n",
    "    assert kernel_size is not None or ks is not None, 'you need to pass a ks'\n",
    "    kernel_size = kernel_size or ks\n",
    "    if padding == 'same':\n",
    "        if kernel_size%2==1:\n",
    "            conv = nn.Conv1d(ni, nf, kernel_size, stride=stride, padding=kernel_size//2 * dilation, dilation=dilation, **kwargs)\n",
    "        else:\n",
    "            conv = SameConv1d(ni, nf, kernel_size, stride=stride, dilation=dilation, **kwargs)\n",
    "    elif padding == 'causal': conv = CausalConv1d(ni, nf, kernel_size, stride=stride, dilation=dilation, **kwargs)\n",
    "    elif padding == 'valid': conv = nn.Conv1d(ni, nf, kernel_size, stride=stride, padding=0, dilation=dilation, **kwargs)\n",
    "    else: conv = nn.Conv1d(ni, nf, kernel_size, stride=stride, padding=padding, dilation=dilation, **kwargs)\n",
    "    init_linear(conv, None, init=init, bias_std=bias_std)\n",
    "    return conv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bs = 2\n",
    "c_in = 3\n",
    "c_out = 5\n",
    "seq_len = 512\n",
    "t = torch.rand(bs, c_in, seq_len)\n",
    "dilation = 1\n",
    "test_eq(CausalConv1d(c_in, c_out, ks=3, dilation=dilation)(t).shape, Conv1d(c_in, c_out, ks=3, padding=\"same\", dilation=dilation)(t).shape)\n",
    "dilation = 2\n",
    "test_eq(CausalConv1d(c_in, c_out, ks=3, dilation=dilation)(t).shape, Conv1d(c_in, c_out, ks=3, padding=\"same\", dilation=dilation)(t).shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bs = 2\n",
    "ni = 3\n",
    "nf = 5\n",
    "seq_len = 6\n",
    "ks = 3\n",
    "t = torch.rand(bs, c_in, seq_len)\n",
    "test_eq(Conv1d(ni, nf, ks, padding=0)(t).shape, (bs, c_out, seq_len - (2 * (ks//2))))\n",
    "test_eq(Conv1d(ni, nf, ks, padding='valid')(t).shape, (bs, c_out, seq_len - (2 * (ks//2))))\n",
    "test_eq(Conv1d(ni, nf, ks, padding='same')(t).shape, (bs, c_out, seq_len))\n",
    "test_eq(Conv1d(ni, nf, ks, padding='causal')(t).shape, (bs, c_out, seq_len))\n",
    "test_error('use kernel_size or ks but not both simultaneously', Conv1d, ni, nf, kernel_size=3, ks=3)\n",
    "test_error('you need to pass a ks', Conv1d, ni, nf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Conv1d(3, 5, kernel_size=(3,), stride=(1,), padding=(1,))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conv = Conv1d(ni, nf, ks, padding='same')\n",
    "init_linear(conv, None, init='auto', bias_std=.01)\n",
    "conv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "CausalConv1d(3, 5, kernel_size=(3,), stride=(1,))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conv = Conv1d(ni, nf, ks, padding='causal')\n",
    "init_linear(conv, None, init='auto', bias_std=.01)\n",
    "conv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Conv1d(3, 5, kernel_size=(3,), stride=(1,))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conv = Conv1d(ni, nf, ks, padding='valid')\n",
    "init_linear(conv, None, init='auto', bias_std=.01)\n",
    "weight_norm(conv)\n",
    "conv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Conv1d(3, 5, kernel_size=(3,), stride=(1,))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conv = Conv1d(ni, nf, ks, padding=0)\n",
    "init_linear(conv, None, init='auto', bias_std=.01)\n",
    "weight_norm(conv)\n",
    "conv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class SeparableConv1d(Module):\n",
    "    def __init__(self, ni, nf, ks, stride=1, padding='same', dilation=1, bias=True, bias_std=0.01):\n",
    "        self.depthwise_conv = Conv1d(ni, ni, ks, stride=stride, padding=padding, dilation=dilation, groups=ni, bias=bias)\n",
    "        self.pointwise_conv = nn.Conv1d(ni, nf, 1, stride=1, padding=0, dilation=1, groups=1, bias=bias)\n",
    "        if bias:\n",
    "            if bias_std != 0:\n",
    "                normal_(self.depthwise_conv.bias, 0, bias_std)\n",
    "                normal_(self.pointwise_conv.bias, 0, bias_std)\n",
    "            else:\n",
    "                self.depthwise_conv.bias.data.zero_()\n",
    "                self.pointwise_conv.bias.data.zero_()\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.depthwise_conv(x)\n",
    "        x = self.pointwise_conv(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bs = 64\n",
    "c_in = 6\n",
    "c_out = 5\n",
    "seq_len = 512\n",
    "t = torch.rand(bs, c_in, seq_len)\n",
    "test_eq(SeparableConv1d(c_in, c_out, 3)(t).shape, (bs, c_out, seq_len))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class AddCoords1d(Module):\n",
    "    \"\"\"Add coordinates to ease position identification without modifying mean and std\"\"\"\n",
    "    def forward(self, x):\n",
    "        bs, _, seq_len = x.shape\n",
    "        cc = torch.linspace(-1,1,x.shape[-1], device=x.device).repeat(bs, 1, 1)\n",
    "        cc = (cc - cc.mean()) / cc.std()\n",
    "        x = torch.cat([x, cc], dim=1)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bs = 2\n",
    "c_in = 3\n",
    "c_out = 5\n",
    "seq_len = 50\n",
    "\n",
    "t = torch.rand(bs, c_in, seq_len)\n",
    "t = (t - t.mean()) / t.std()\n",
    "test_eq(AddCoords1d()(t).shape, (bs, c_in + 1, seq_len))\n",
    "new_t = AddCoords1d()(t)\n",
    "test_close(new_t.mean(),0, 1e-2)\n",
    "test_close(new_t.std(), 1, 1e-2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class ConvBlock(nn.Sequential):\n",
    "    \"Create a sequence of conv1d (`ni` to `nf`), activation (if `act_cls`) and `norm_type` layers.\"\n",
    "    def __init__(self, ni, nf, kernel_size=None, ks=3, stride=1, padding='same', bias=None, bias_std=0.01, norm='Batch', zero_norm=False, bn_1st=True,\n",
    "                 act=nn.ReLU, act_kwargs={}, init='auto', dropout=0., xtra=None, coord=False, separable=False,  **kwargs):\n",
    "        kernel_size = kernel_size or ks\n",
    "        ndim = 1\n",
    "        layers = [AddCoords1d()] if coord else []\n",
    "        norm_type = getattr(NormType,f\"{snake2camel(norm)}{'Zero' if zero_norm else ''}\") if norm is not None else None\n",
    "        bn = norm_type in (NormType.Batch, NormType.BatchZero)\n",
    "        inn = norm_type in (NormType.Instance, NormType.InstanceZero)\n",
    "        if bias is None: bias = not (bn or inn)\n",
    "        if separable: conv = SeparableConv1d(ni + coord, nf, ks=kernel_size, bias=bias, stride=stride, padding=padding, **kwargs)\n",
    "        else: conv = Conv1d(ni + coord, nf, ks=kernel_size, bias=bias, stride=stride, padding=padding, **kwargs)\n",
    "        act = None if act is None else act(**act_kwargs)\n",
    "        if not separable: init_linear(conv, act, init=init, bias_std=bias_std)\n",
    "        if   norm_type==NormType.Weight:   conv = weight_norm(conv)\n",
    "        elif norm_type==NormType.Spectral: conv = spectral_norm(conv)\n",
    "        layers += [conv]\n",
    "        act_bn = []\n",
    "        if act is not None: act_bn.append(act)\n",
    "        if bn: act_bn.append(BatchNorm(nf, norm_type=norm_type, ndim=ndim))\n",
    "        if inn: act_bn.append(InstanceNorm(nf, norm_type=norm_type, ndim=ndim))\n",
    "        if bn_1st: act_bn.reverse()\n",
    "        if dropout: layers += [nn.Dropout(dropout)]\n",
    "        layers += act_bn\n",
    "        if xtra: layers.append(xtra)\n",
    "        super().__init__(*layers)\n",
    "\n",
    "Conv = named_partial('Conv', ConvBlock, norm=None, act=None)\n",
    "ConvBN = named_partial('ConvBN', ConvBlock, norm='Batch', act=None)\n",
    "CoordConv = named_partial('CoordConv', ConvBlock, norm=None, act=None, coord=True)\n",
    "SepConv = named_partial('SepConv', ConvBlock, norm=None, act=None, separable=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class ResBlock1dPlus(Module):\n",
    "    \"Resnet block from `ni` to `nh` with `stride`\"\n",
    "#     @delegates(ConvLayer.__init__)\n",
    "    def __init__(self, expansion, ni, nf, coord=False, stride=1, groups=1, reduction=None, nh1=None, nh2=None, dw=False, g2=1,\n",
    "                 sa=False, sym=False, norm='Batch', zero_norm=True, act_cls=defaults.activation, ks=3,\n",
    "                 pool=AvgPool, pool_first=True, **kwargs):\n",
    "        if nh2 is None: nh2 = nf\n",
    "        if nh1 is None: nh1 = nh2\n",
    "        nf,ni = nf*expansion,ni*expansion\n",
    "        k0 = dict(norm=norm, zero_norm=False, act=act_cls, **kwargs)\n",
    "        k1 = dict(norm=norm, zero_norm=zero_norm, act=None, **kwargs)\n",
    "        convpath  = [ConvBlock(ni,  nh2, ks, coord=coord, stride=stride, groups=ni if dw else groups, **k0),\n",
    "                     ConvBlock(nh2,  nf, ks, coord=coord, groups=g2, **k1)\n",
    "        ] if expansion == 1 else [\n",
    "                     ConvBlock(ni,  nh1, 1, coord=coord, **k0),\n",
    "                     ConvBlock(nh1, nh2, ks, coord=coord, stride=stride, groups=nh1 if dw else groups, **k0),\n",
    "                     ConvBlock(nh2,  nf, 1, coord=coord, groups=g2, **k1)]\n",
    "        if reduction: convpath.append(SEModule(nf, reduction=reduction, act_cls=act_cls))\n",
    "        if sa: convpath.append(SimpleSelfAttention(nf,ks=1,sym=sym))\n",
    "        self.convpath = nn.Sequential(*convpath)\n",
    "        idpath = []\n",
    "        if ni!=nf: idpath.append(ConvBlock(ni, nf, 1, coord=coord, act=None, **kwargs))\n",
    "        if stride!=1: idpath.insert((1,0)[pool_first], pool(stride, ndim=1, ceil_mode=True))\n",
    "        self.idpath = nn.Sequential(*idpath)\n",
    "        self.act = defaults.activation(inplace=True) if act_cls is defaults.activation else act_cls()\n",
    "\n",
    "    def forward(self, x): return self.act(self.convpath(x) + self.idpath(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def SEModule1d(ni, reduction=16, act=nn.ReLU, act_kwargs={}):\n",
    "    \"Squeeze and excitation module for 1d\"\n",
    "    nf = math.ceil(ni//reduction/8)*8\n",
    "    assert nf != 0, 'nf cannot be 0'\n",
    "    return SequentialEx(nn.AdaptiveAvgPool1d(1),\n",
    "                        ConvBlock(ni, nf, ks=1, norm=None, act=act, act_kwargs=act_kwargs),\n",
    "                        ConvBlock(nf, ni, ks=1, norm=None, act=nn.Sigmoid), ProdLayer())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = torch.rand(8, 32, 12)\n",
    "test_eq(SEModule1d(t.shape[1], 16, act=nn.ReLU, act_kwargs={})(t).shape, t.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def Norm(nf, ndim=1, norm='Batch', zero_norm=False, init=True, **kwargs):\n",
    "    \"Norm layer with `nf` features and `ndim` with auto init.\"\n",
    "    assert 1 <= ndim <= 3\n",
    "    nl = getattr(nn, f\"{snake2camel(norm)}Norm{ndim}d\")(nf, **kwargs)\n",
    "    if nl.affine and init:\n",
    "        nl.bias.data.fill_(1e-3)\n",
    "        nl.weight.data.fill_(0. if zero_norm else 1.)\n",
    "    return nl\n",
    "\n",
    "BN1d = partial(Norm, ndim=1, norm='Batch')\n",
    "IN1d = partial(Norm, ndim=1, norm='Instance')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bs = 2\n",
    "ni = 3\n",
    "nf = 5\n",
    "sl = 4\n",
    "ks = 5\n",
    "\n",
    "t = torch.rand(bs, ni, sl)\n",
    "test_eq(ConvBlock(ni, nf, ks)(t).shape, (bs, nf, sl))\n",
    "test_eq(ConvBlock(ni, nf, ks, padding='causal')(t).shape, (bs, nf, sl))\n",
    "test_eq(ConvBlock(ni, nf, ks, coord=True)(t).shape, (bs, nf, sl))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(BN1d(ni)(t).shape, (bs, ni, sl))\n",
    "test_eq(BN1d(ni).weight.data.mean().item(), 1.)\n",
    "test_eq(BN1d(ni, zero_norm=True).weight.data.mean().item(), 0.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ConvBlock(\n",
       "  (0): AddCoords1d()\n",
       "  (1): Conv1d(4, 5, kernel_size=(5,), stride=(1,), padding=(2,), bias=False)\n",
       "  (2): BatchNorm1d(5, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (3): SiLU()\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_eq(ConvBlock(ni, nf, ks, norm='batch', zero_norm=True)[1].weight.data.unique().item(), 0)\n",
    "test_ne(ConvBlock(ni, nf, ks, norm='batch', zero_norm=False)[1].weight.data.unique().item(), 0)\n",
    "test_eq(ConvBlock(ni, nf, ks, bias=False)[0].bias, None)\n",
    "ConvBlock(ni, nf, ks, act=Swish, coord=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class LinLnDrop(nn.Sequential):\n",
    "    \"Module grouping `LayerNorm1d`, `Dropout` and `Linear` layers\"\n",
    "    def __init__(self, n_in, n_out, ln=True, p=0., act=None, lin_first=False):\n",
    "        layers = [nn.LayerNorm(n_out if lin_first else n_in)] if ln else []\n",
    "        if p != 0: layers.append(nn.Dropout(p))\n",
    "        lin = [nn.Linear(n_in, n_out, bias=not ln)]\n",
    "        if act is not None: lin.append(act)\n",
    "        layers = lin+layers if lin_first else layers+lin\n",
    "        super().__init__(*layers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "LinLnDrop(\n",
       "  (0): LayerNorm((2,), eps=1e-05, elementwise_affine=True)\n",
       "  (1): Dropout(p=0.5, inplace=False)\n",
       "  (2): Linear(in_features=2, out_features=3, bias=False)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "LinLnDrop(2, 3, p=.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class LambdaPlus(Module):\n",
    "    def __init__(self, func, *args, **kwargs): self.func,self.args,self.kwargs=func,args,kwargs\n",
    "    def forward(self, x): return self.func(x, *self.args, **self.kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class Squeeze(Module):\n",
    "    def __init__(self, dim=-1): self.dim = dim\n",
    "    def forward(self, x): return x.squeeze(dim=self.dim)\n",
    "    def __repr__(self): return f'{self.__class__.__name__}(dim={self.dim})'\n",
    "\n",
    "\n",
    "class Unsqueeze(Module):\n",
    "    def __init__(self, dim=-1): self.dim = dim\n",
    "    def forward(self, x): return x.unsqueeze(dim=self.dim)\n",
    "    def __repr__(self): return f'{self.__class__.__name__}(dim={self.dim})'\n",
    "\n",
    "\n",
    "class Add(Module):\n",
    "    def forward(self, x, y): return x.add(y)\n",
    "    def __repr__(self): return f'{self.__class__.__name__}'\n",
    "\n",
    "\n",
    "class Concat(Module):\n",
    "    def __init__(self, dim=1): self.dim = dim\n",
    "    def forward(self, *x): return torch.cat(*x, dim=self.dim)\n",
    "    def __repr__(self): return f'{self.__class__.__name__}(dim={self.dim})'\n",
    "\n",
    "\n",
    "class Unfold(Module):\n",
    "    def __init__(self, dim, size, step=1): self.dim, self.size, self.step =  dim, size, step\n",
    "    def forward(self, x:Tensor) -> Tensor: return x.unfold(dimension=self.dim, size=self.size, step=self.step)\n",
    "    def __repr__(self): return f\"{self.__class__.__name__}(dim={self.dim}, size={self.size}, step={self.step})\"\n",
    "\n",
    "\n",
    "class Permute(Module):\n",
    "    def __init__(self, *dims): self.dims = dims\n",
    "    def forward(self, x:Tensor) -> Tensor: return x.permute(self.dims)\n",
    "    def __repr__(self): return f\"{self.__class__.__name__}(dims={', '.join([str(d) for d in self.dims])})\"\n",
    "\n",
    "\n",
    "class Transpose(Module):\n",
    "    def __init__(self, *dims, contiguous=True): self.dims, self.contiguous = dims, contiguous\n",
    "    def forward(self, x):\n",
    "        x = x.contiguous()\n",
    "        if self.contiguous: return x.transpose(*self.dims).contiguous()\n",
    "        else: return x.transpose(*self.dims)\n",
    "    def __repr__(self):\n",
    "        if self.contiguous: return f\"{self.__class__.__name__}(dims={', '.join([str(d) for d in self.dims])}).contiguous()\"\n",
    "        else: return f\"{self.__class__.__name__}({', '.join([str(d) for d in self.dims])})\"\n",
    "\n",
    "\n",
    "class View(Module):\n",
    "    def __init__(self, *shape): self.shape = shape\n",
    "    def forward(self, x):\n",
    "        return x.view(x.shape[0], -1).contiguous() if not self.shape else x.view(-1).contiguous() if self.shape == (-1,) else \\\n",
    "            x.view(x.shape[0], *self.shape).contiguous()\n",
    "    def __repr__(self): return f\"{self.__class__.__name__}({', '.join(['bs'] + [str(s) for s in self.shape])})\"\n",
    "\n",
    "\n",
    "class Reshape(Module):\n",
    "    def __init__(self, *shape): self.shape = shape\n",
    "    def forward(self, x):\n",
    "        return x.contiguous().reshape(x.shape[0], -1) if not self.shape else x.contiguous().reshape(-1) if self.shape == (-1,) else x.contiguous().reshape(x.shape[0], *self.shape)\n",
    "    def __repr__(self): return f\"{self.__class__.__name__}({', '.join(['bs'] + [str(s) for s in self.shape])})\"\n",
    "\n",
    "\n",
    "class Max(Module):\n",
    "    def __init__(self, dim=None, keepdim=False): self.dim, self.keepdim = dim, keepdim\n",
    "    def forward(self, x): return x.max(self.dim, keepdim=self.keepdim)[0]\n",
    "    def __repr__(self): return f'{self.__class__.__name__}(dim={self.dim}, keepdim={self.keepdim})'\n",
    "\n",
    "\n",
    "class LastStep(Module):\n",
    "    def forward(self, x): return x[..., -1]\n",
    "    def __repr__(self): return f'{self.__class__.__name__}()'\n",
    "\n",
    "\n",
    "class SoftMax(Module):\n",
    "    \"SoftMax layer\"\n",
    "    def __init__(self, dim=-1):\n",
    "        self.dim = dim\n",
    "    def forward(self, x):\n",
    "        return F.softmax(x, dim=self.dim)\n",
    "    def __repr__(self): return f'{self.__class__.__name__}(dim={self.dim})'\n",
    "\n",
    "\n",
    "class Clamp(Module):\n",
    "    def __init__(self, min=None, max=None):\n",
    "        self.min, self.max = min, max\n",
    "    def forward(self, x):\n",
    "        return x.clamp(min=self.min, max=self.max)\n",
    "    def __repr__(self): return f'{self.__class__.__name__}(min={self.min}, max={self.max})'\n",
    "\n",
    "\n",
    "class Clip(Module):\n",
    "    def __init__(self, min=None, max=None):\n",
    "        self.min, self.max = min, max\n",
    "\n",
    "    def forward(self, x):\n",
    "        if self.min is not None:\n",
    "            x = torch.maximum(x, self.min)\n",
    "        if self.max is not None:\n",
    "            x = torch.minimum(x, self.max)\n",
    "        return x\n",
    "    def __repr__(self): return f'{self.__class__.__name__}()'\n",
    "\n",
    "\n",
    "class ReZero(Module):\n",
    "    def __init__(self, module):\n",
    "        self.module = module\n",
    "        self.alpha = nn.Parameter(torch.zeros(1))\n",
    "    def forward(self, x):\n",
    "        return x + self.alpha * self.module(x)\n",
    "\n",
    "\n",
    "Noop = nn.Sequential()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Transpose(dims=1, 2).contiguous(),\n",
       " Permute(dims=0, 2, 1),\n",
       " View(bs, -1, 2, 10),\n",
       " Transpose(dims=1, 2).contiguous(),\n",
       " Reshape(bs, -1, 2, 10),\n",
       " Sequential())"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 2\n",
    "nf = 5\n",
    "sl = 4\n",
    "\n",
    "t = torch.rand(bs, nf, sl)\n",
    "test_eq(Permute(0,2,1)(t).shape, (bs, sl, nf))\n",
    "test_eq(Max(1)(t).shape, (bs, sl))\n",
    "test_eq(Transpose(1,2)(t).shape, (bs, sl, nf))\n",
    "test_eq(Transpose(1,2, contiguous=True)(t).shape, (bs, sl, nf))\n",
    "test_eq(View(-1, 2, 10)(t).shape, (bs, 1, 2, 10))\n",
    "test_eq(Reshape(-1, 2, 10)(t).shape, (bs, 1, 2, 10))\n",
    "test_eq(Reshape()(t).shape, (2, 20))\n",
    "test_eq(Reshape(-1)(t).shape, (40,))\n",
    "Transpose(1,2), Permute(0,2,1), View(-1, 2, 10), Transpose(1,2, contiguous=True), Reshape(-1, 2, 10), Noop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class DropPath(nn.Module):\n",
    "    \"\"\"Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).\n",
    "\n",
    "    It's similar to Dropout but it drops individual connections instead of nodes.\n",
    "    Original code in https://github.com/rwightman/pytorch-image-models (timm library)\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, p=None):\n",
    "        super().__init__()\n",
    "        self.p = p\n",
    "\n",
    "    def forward(self, x):\n",
    "        if self.p == 0. or not self.training: return x\n",
    "        keep_prob = 1 - self.p\n",
    "        shape = (x.shape[0],) + (1,) * (x.ndim - 1)\n",
    "        random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)\n",
    "        random_tensor.floor_()\n",
    "        output = x.div(keep_prob) * random_tensor\n",
    "#         output = x.div(random_tensor.mean()) * random_tensor # divide by the actual mean to mantain the input mean?\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = torch.ones(100,2,3)\n",
    "test_eq(DropPath(0.)(t), t)\n",
    "assert DropPath(0.5)(t).max() >= 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class Sharpen(Module):\n",
    "    \"This is used to increase confidence in predictions - MixMatch paper\"\n",
    "    def __init__(self, T=.5): self.T = T\n",
    "    def forward(self, x):\n",
    "        x = x**(1. / self.T)\n",
    "        return x / x.sum(dim=1, keepdims=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAiMAAAGdCAYAAADAAnMpAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjAsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvlHJYcgAAAAlwSFlzAAAPYQAAD2EBqD+naQAARgBJREFUeJzt3Qm8TPX/x/H3dbmWZEv2ndIihEhaf5RKi7afpIgoKhUtUiJt2pQW0aa9HyppsbSQSpSQSlkS4S+77Fzb/B+f8zV30b3cZWbOLK/n43E63zlzZs73nube+fgun29SIBAICAAAwCcF/LowAACAIRgBAAC+IhgBAAC+IhgBAAC+IhgBAAC+IhgBAAC+IhgBAAC+IhgBAAC+KqgYsG/fPv399986/PDDlZSU5Hd1AABADlhe1S1btqhSpUoqUKBAbAcjFohUrVrV72oAAIA8WL58uapUqRLbwYi1iAR/mBIlSvhdHQAAkAObN2/2GhOC3+MxHYwEu2YsECEYAQAgthxqiAUDWAEAgK8IRgAAgK8IRgAAgK8IRgAAgK8IRgAAgK8IRgAAgK8IRgAAgK8IRgAAgK8IRgAAQGwFI998840uvPBCb9Eby6g2duzYQ75mypQpatSokQoXLqw6dero9ddfz2t9AQBAogcj27ZtU4MGDTR06NAcnb9kyRK1adNGZ511lubMmaPbbrtNXbt21WeffZaX+gIAgDiT67VpzjvvPG/LqeHDh6tmzZoaPHiw9/jYY4/V1KlT9fTTT6t169a5vTwAAIgzYV8ob/r06WrVqlWmYxaEWAtJdlJTU70t46p/AABEu0BA2rFD2r49682e273bbXv2pJd379/27Tv4tnfvv4/ZNTNe/8D6ZPlcaqoCm7dIW7bYl6y0eZN6jaivGg1LKS6DkVWrVql8+fKZjtljCzB27NihokWL/us1gwYN0sCBA8NdNQAADsq+7BctkhYskFavltatk9avlzZskP75x21W3rhR2rTJfa8fGBBEp8L7t7JpR9rPnBu/wUhe9O3bV7179057bIFL1apVfa0TACA+WWvD3LnSL79Iy5ZJy5dLf/0lLV3qNmvNyIvChaXDDpOKFUvf7FhKilSwoFSoUOatYEEpOdltBQrkbEtKcnsvalq1SlqzWkl//y2tWSNtWO81vyQpPTrKWFbRokoqUUKyrXQpVapTW34JezBSoUIFrbZwMgN7XKJEiSxbRYzNurENAIBQsq4R+57+7jtpxgxp1ixp9mzXqpEd+6qqVUuqWVM64gipbFmpdGm3lSnj9qVKua1kSal4cfcaCyrCZvt26YcfpG+/laZOtTER0tat/z7PvkuPOUY67ji3HXusZP+4P/54FylFibAHI82bN9f48eMzHfviiy+84wAAhLvVwwKODz+U3nvPZni6RoQDWePAiSe6oKNKFalGDal6dbevVs21XPhq924XQX3/vQs+Pv/cHTvwhzjlFKlRI6lpUxdwWAQV1qjIp2Bk69atWmQdaBmm7tqU3TJlyqhatWpeF8uKFSv05ptves93795dzz//vO666y516dJFkydP1ujRozVu3LjQ/iQAAEhasUKaPFmylFbTpkk7d2Z+3ro1jj5a+s9/pMaN3Xe3fW/7HnAcaNs2aeJEt9l35sqVmZ+3qOm006RTT3Wb/RAxEHiEJBiZOXOmlzMkKDi2o1OnTl4ys5UrV2qZdbrtZ9N6LfDo1auXnnnmGVWpUkWvvPIK03oBACFjA0jtO/vVV6Uvv8z8nI3RaNtWOuccS09hkyii+Dt71y5p0iTp449dU46Nlg2y/iCbnWqtHvY9bFGUDRqJA0mBQPSP+7UBrCVLltSmTZu8sSYAABjrsXj+eemjjzK3gFiLx/nnS+3bu1aQqA0+zI4d0tdfS++8I33ySeYBLNZPdPHFLoo6/XQ3GCWG5PT7Oypn0wAAkB37rn7wQddzMX9++nEbm2kNBzfc4Hosop7NA7ZIasgQae3a9OPWdHPppdIFF0hnnx2F/UehRzACAIgZ9p19xhnSvHnusbV4XHut1KNHDPVa2LzhJ55wg1psVow58kjXlNOli9SiRZQ35YQewQgAICbY7FX7rrbWEOutePll9/1tQyligiUzuece16QTnNJjrSAPPyx17JgQLSDZIRgBAER9bpB+/aTHHnOPDz9cmjDBNSDEBBuaaQHH/fe7ucbGpvLcdZcbiJqSokRHMAIAiOpAxFo/vvjCPbZhFMOGuVmtMRGEfPqp65Kx5GSmWTM34MXGgiANwQgAICrZmi9NmrhEZZZI9O23pcsvV2ywxWxuv911yQSTmwwY4Jp4vPztyIhgBAAQlTp3doGIsfwhMRGIWEr2m25ykVNwXEjXrtItt0gnnOB37aIWwQgAIOryft17r8v7ZWxoRYcOin6Wqv2//3Ur7RnLjvrQQy4/CA6KYAQAEFUNC/bd/dNP7nGbNukDV6PWunVumo8lLDM2oGXECMaF5AIdVwCAqGErjAQDkf79pbFjFd1sFT7LsBYMRC680P0ABCK5QssIACBqejleecWVX3xRuv56RS9bMPbWW6XPPkufrmuL47DuWp7QMgIAiIp8YGee6WbD2qr31usRtUaPdovfjB/vApFLLpHWrCEQyQeCEQCA77lEbIBqaqp00kluzbiC0dpu/8wzUrt2bl0Zyz//88/SmDEunTvyjGAEAOD7WjO//OIeDx8uVa2q6PTCC9Jtt7nyNddIM2ZI9ev7Xau4QDACAPCNpd+YNk0qXlwaOdI1NkQlawmxEbXG+pMs8UmCLWYXTgQjAABfzJnjAhDzzjuu9yMqrVolXXGFtH69VK2ay02fwIvahUO09soBAOLcs8+6fb16bkZsVLJpuuee6wao2oJ2r70WxQNaYhctIwCAiLIs6dbjYd/r5vHHpaQkRRdrBbEF7po3d4GIpXL/8Ue32i5CjvAOABAxNnXXMqZ/8IF7bMu4nHeeose2bS7lqwUiO3e6Yyef7HKIlCzpd+3iFsEIACBiZs9OD0SsZaRTJ0VXLvqWLd0sGWMJT3r2lHr0kIoU8bt2cY1gBAAQsSVcrrzSlZs2la69VtGVUbVVK2npUje1xxa4s6k+Udd/FJ8IRgAAEXHKKe47v1gx6cknFV1Rks0p3rJFOuwwt86MTd9FxBCMAADCzsaA/vGHK9v6M6edpujx1FMuEAkukGPTexBRzKYBAIR90KoNVDU1aqR31UTNojg2YNX07Usg4hOCEQBAWN1zj/T++678+utRNAzD1pWxKbs217hKFWngQL9rlLAIRgAAYbN4sZsla4YNc+vQREVTjWVcsym7QZ9+SlZVHzFmBAAQNjaNd+9eN3i1e3e/a7M/EOncWXrjDfe4XDmXQ6RBA79rltBoGQEAhMWePek9HxdfHCUV6tIlPRC57z5p2TLpxBP9rlnCo2UEABAWU6a4hKYmKpKbXXCB9Nlnrvz88+mjauE7WkYAAGGZymsL3ZqOHaXy5X2szOrV0uWXpwcir75KIBJlCEYAACG1YIHUsKG0caN77PtYkT590nPQ2zRe66pBVCEYAQCE1NVXSytXukyr333nFr71xaZNrnkmOEbkxhulu+7yqTI4GMaMAABCZswYaeZMl0vkp5+ko4/2qSIWDVWqlP64TRu33gyiEi0jAICQsIkpwcXvrHXEt0DEkptkbI55912XR6R0aZ8qhEOhZQQAEBJjx7olXqpXl154wcc8IjZi1lbfrVBBmjDBDWBBVKNlBAAQsqm8pkcPqXhxHyqwc6d06aVuoEowOiIQiQm0jAAA8s2Wd/n6a1c+80yfWkTatZM+/tgNWBk+XGrWzIeKIC8IRgAAIVn8dsMG1yLSqJEPgUirVtLkyVKBAtLnn0stW0a4EsgPumkAAPk2erTbn3qqD+vN2RQeC0TMoEEEIjGIYAQAkC8WBzz8sCuffbYPrSK2xozp1o08IjGKYAQAkC+33OL2NnP25psjfPGhQ6V586SUFOmRRyJ8cYQKwQgAIM/++Uf67TdX/uQTFxNEzK5dUs+ernzrrVLZshG8OEKJYAQAkGcDBrh9uXJSkyYRvviwYenl/v0jfHGEEsEIACDPLPW7saEahQtHeC7xc8/5nNgEoUIwAgDIs4UL3d5m1kZU167Sn3+6IMRW4kVMIxgBAOTJN99I69e7HGN16kTwwrNmSa+95soWiBx+eAQvjnAgGAEA5MmTT7p9hw7SYYdF6KI2c+ayy1z5v/+VbrwxQhdGOJGBFQCQa3v2SNOnp6f3iIg1a1yKd1uNr1QpaeDACF0Y4UYwAgDIlb173TIw69a5qbwNGkRwJT4LRIzNJ65UKUIXRrjRTQMAyJXx410GdjN4sFSyZIQioKefduVrryUQiTMEIwCAXPngA7e3HpOIZVz98kvp++9d+bbbInRRRArBCAAgxzZulN55x5U7d47ghd94w+0vvzyC/UKIFIIRAECODRniBq9Wrixdf32ELvr889LIka58ww0RuigiiWAEAJAjw4dLDz6YPq3X8ouEnV3I1p+x1Xkt02rLlhG4KCKNYAQAcEg2icViAcvCfuWVbjZN2P3xh9Svnys/8IBboTciERAijWAEAJDjNWjMm29GKCZ49lkpNVU67TQXlBCIxC2CEQBAjiazmHPOkQoVitBFJ0xInz1DIBLX8hSMDB06VDVq1FCRIkXUrFkzzZgx46DnDxkyRHXr1lXRokVVtWpV9erVSzt37sxrnQEAEbR4sfTII658ySURuui4cW4hPHP66RG6KGImGBk1apR69+6tAQMGaPbs2WrQoIFat26tNZamNwvvvvuu7r77bu/8efPm6dVXX/Xe45577glF/QEAYRZck+6IIyKU+t1W3wvOG771Vqls2QhcFDEVjDz11FPq1q2bOnfurOOOO07Dhw9XsWLFNGLEiCzPnzZtmlq0aKGrrrrKa00555xz1L59+0O2pgAA/GcDVt9+25VvuklKTo7ARW+/XVq7VjrmGGnQoAhcEDEVjOzatUuzZs1Sq1at0t+gQAHv8fTgikkHOOWUU7zXBIOPxYsXa/z48Tr//POzvU5qaqo2b96caQMARN5LL0l//ZUeI4TdokXSW2+lT+stWjQCF0VMLZS3bt067d27V+XLl8903B7Pnz8/y9dYi4i97tRTT1UgENCePXvUvXv3g3bTDBo0SANZjREAfO+esem85pprpBIlwnxBW3nvzDNdc4xlWT3IP1oRX8I+m2bKlCl65JFH9MILL3hjTMaMGaNx48bpwWDmnCz07dtXmzZtStuWL18e7moCAA4Q/Ddj9equhSTsevWSVqyQqlSRPvyQGTQJJFctI2XLllVycrJWr16d6bg9rlChQpavue+++3TNNdeoa9eu3uMTTjhB27Zt0/XXX697773X6+Y5UOHChb0NAOCPWbOkVatc+dNPpSJFwnzB7dttxoMrv/++VLNmmC+ImG0ZSUlJUePGjTVp0qS0Y/v27fMeN2/ePMvXbN++/V8BhwU0xrptAADRxf40B1fjtWyr9epFaCqvdc8ceaRbDhgJJVctI8am9Xbq1ElNmjRR06ZNvRwi1tJhs2tMx44dVblyZW/ch7nwwgu9GTgnnniil5Nk0aJFXmuJHQ8GJQCA6GEL5H7/vbVSS48+GqGL2mBVU7t2hC6ImA5G2rVrp7Vr16p///5atWqVGjZsqIkTJ6YNal22bFmmlpB+/fopKSnJ269YsUJHHnmkF4g8/PDDof1JAAAhYb0k5qqr3HiRsNu0KT3ffMSWAkY0SQrEQF+JTe0tWbKkN5i1RNiHcwNA4rJMCiVLuvIvv9g4vwg1xVx7rRsnYllXGbgaN3L6/c3aNACANF26uH2tWhEaK2I++CB9/jCBSEIiGAEAZJpFY268MUJxgaV+Dy6Id8opEbggohHBCAAgbUG8YLbVCy6I0EX79ZP27JGOOko6++wIXRTRhmAEAOAZNsztGzaU6tSJwAW3bElf+Obxx219kQhcFNGI//MAAC/Fx5tvpq9BE5HMC3feKW3dKtWtK118cQQuiGhFMAIA0DPPSGvWSMWKSZdfHoELbtwovfhiehp4Bq4mNIIRAIC+/NLtL7ooAqnfzbRpbm9dM+QWSXgEIwAALVjg9t27R+BiNmB1f9ZudexIqwgIRgAg0e3cKS1Z4so2fCPsbH0z6xMyPXpE4IKIdgQjAJDgFi1yA1gt8+r+lT3CK7g6rwUiTZtG4IKIdgQjAJDg5s9PbxUJe4+J9QcFp+20bx/miyFWEIwAQIKbN8/tjz02gslMrEWkRYsIXBCxgGAEABLcb79FKBjZtUv6/PP0kbIkOcN+fBIAIIHt2CGNH+/Kxx8fxgutXCmdfLJrhilaVDr33DBeDLGGYAQAEtjXX7us7BUqSK1bh3H2TM2a0k8/SSkpbsxIxYphuhhiUUG/KwAA8M+UKW5vjRaFCoXhArbyniU1S011Ach77zFWBP9CMAIACez1192+Vq0wvHkg4JpbbDng0qVdy0hE5g4j1tBNAwAJyhKhrl/vymFZp+7PP6WFC135m28IRJAtghEASFD/938uILFhHGHpOfn+e7dv3lyqVy8MF0C8IBgBgATOvGpsbGlychgXw7MBKcBBEIwAQIJ6+GG3b9w4DG8+d640fLgrn3FGGC6AeEIwAgAJaN269Jk0nTqF+M23b3frztgA1kaNpAsvDPEFEG8IRgAgAb3/vtsfeaR09tkhfvNevaSpU135qafItIpD4hMCAAno99/dvmPHEC+OZ5lWX3rJle++my4a5AjBCAAkGJtB89FHrtygQYjf/NtvM7eQADlAMAIACWb+fGnZMumww6TLLw/xm0+Y4PbduknlyoX4zRGvCEYAIMEsXer2deq4NetC5p9/pJEjXfnaa0P4xoh3BCMAkGCsVcRUrx7CN7WZM/ffL+3cKdWv7xKdATlEMAIACdoyEtJg5I47pGefdeXu3UM8KhbxjmAEABK0ZaRatRDOoLEpvMHpOTZeBMgFVu0FgAQT0paRffuk3r1duWRJ6Y03QvCmSDS0jABAArHYwWbTmBo1QvCG77yTPmiVqbzII4IRAEggP/4obdgglSghNWyYzzfbvFkaONCVbY7wgAGhqCISEMEIACQQm/Bizj1XKlQon292223Sn39KBQtKffqEonpIUAQjAJAgLEv7xImufOedIRgF+9prrmxv2qRJvuuHxEUwAgAJ4n//c/sOHfIZO9jAk/POc+VjjpFatgxJ/ZC4CEYAIAFYctRvvnHlG27I55tdf336SnvBKb1APhCMAEAC+Pln16BRpox06qn5eKMHHpBefTW93yfYQgLkA8EIACSAJUvc/qST8pEc9eWX02fMkNwMIUQwAgAJFIzkObfIX3+5NO/m0kulESNCVjeAYAQA4tzevdKoUa58wgl5fJPXX3f9PJacZPRoKTk5lFVEgiMYAYA49+uv0sKF0uGHu96VPEUzwWm8d91FIIKQIxgBgDj3ySduX7u2C0hy7YUXXF6RUqWkSy4JdfUAghEAiHczZ7r9EUfk4cWBgHTLLa7ctatUpEhI6wYYghEAiHPBlCB33JGHF3/7beZpvUAYEIwAQBzbsUNavNiVc70wng1YDa7EW6GCVLRoyOsHGIIRAIhjv/3mYgrroilfPpcvHjNGmj3blceNC0f1AA/BCADEsRkz0ltFcpXszCKY3r1d2VK2NmoUlvoBhmAEAOI8Dbxp3jyXLxw8WFq+3JUffzzk9QIyIhgBgDgWjCdylXl1yhSXT8TYTJpcRzJA7hCMAEAcsyzupmrVXLxowoT08pAhIa8TcCCCEQCIU5Y41TKvmmOPzeGL9uxx6d6Nrc6b51X1gJwjGAGAOLV2rQtIChSQKlXK4YsmTnTNKUceKbVvH+YaAg7BCADEqVWr3L5cuVwsJ/Pjj25/zjnkFUHEEIwAQJyPF6lYMYcv2LpVevRRVz7++LDVCzgQwQgAxKngeJHjjsvhC3r2lHbtcs0odNEg2oORoUOHqkaNGipSpIiaNWumGcGsOtnYuHGjbrrpJlWsWFGFCxfW0UcfrfHjx+e1zgCAHNiyxe1tsd1DWrpUGjXKlQcMyOVcYCB/Cub2BaNGjVLv3r01fPhwLxAZMmSIWrdurQULFqicdUweYNeuXTr77LO9595//31VrlxZS5cuVakc/XYAAPJq82a3P/zwQ5yYmiq1aeMWsjnlFKlfv0hUD8h7MPLUU0+pW7du6ty5s/fYgpJx48ZpxIgRuvvuu/91vh3fsGGDpk2bpkKFCnnHrFUFABBe06e7ffXqOeiesUVsrHvmrbeYzovo7qaxVo5Zs2apVatW6W9QoID3eHrwU3+Ajz/+WM2bN/e6acqXL6969erpkUce0V6bb5aN1NRUbd68OdMGAMi5NWukmTNd+eKLDzH/9+WXXdm6aWrVikj9gDwHI+vWrfOCCAsqMrLHq4JzyA6wePFir3vGXmfjRO677z4NHjxYDz30ULbXGTRokEqWLJm2Vc1V6kAAgC0nEwi4wasHnU3Tt6/b164tXXZZpKoHRHY2zb59+7zxIi+99JIaN26sdu3a6d577/W6d7LTt29fbdq0KW1bHlxcAQCQI9995/a3356DJGc5OhGIkjEjZcuWVXJyslavXp3puD2uUKFClq+xGTQ2VsReF3Tsscd6LSnW7ZOSkvKv19iMG9sAAHnrovn+e1du2vQgJ9pJK1a4FK0dO0aqekD+WkYscLDWjUmTJmVq+bDHNi4kKy1atNCiRYu884IWLlzoBSlZBSIAgPx5+un09WiyzTFifTi9e7ty27bSYYdFrH5AvrtpbFrvyy+/rDfeeEPz5s1Tjx49tG3btrTZNR07dvS6WYLseZtNc+utt3pBiM28sQGsNqAVABBaNmg1uNDuww+7Ro8sTZ7spttYK/Tzz0eyikD+p/bamI+1a9eqf//+XldLw4YNNXHixLRBrcuWLfNm2ATZ4NPPPvtMvXr1Uv369b08IxaY9OnTJ7eXBgAcIsnZVVdJO3e6pWWswSPbVpH+/V25W7dc5IsHwiMpELBPZXSzqb02q8YGs5YoUcLv6gBAVLIZutdfb13q0pIlB1mpd+5c6YQTpIIF3YlVqkS4pkgUm3P4/c3aNAAQJ375xe2tFzzbQMQmIFxyiSs3bkwggqhAMAIAccDmCLz3nis3bHiQk6wfZ9EiyWZAPvNMJKsIhG7MCAAg+ixe7Bo9TMuW2Zz0wANu4KoNWrX8Ig0aRLKKQLZoGQGAOGDpQoz1ulSunM1J777r9jall0AEUYRgBADiwFdfuX39+gfpolm6NH0GDRBFCEYAIMbZnMg330wfk5olW1Zj1y6XeIRBq4gyBCMAEAddNDZD13Ttms1Jb7/t9iedJBUqFLG6ATlBMAIAMe7LL93ehoFUq5bNSa++6vbt2kWsXkBOEYwAQIwbO/YQXTTLlqU3nbRvH7F6ATlFMAIAMSw1VQquXdqlSzYnPfmk29uCptmssA74iWAEAGLYJ59IW7e6GCPLxdM3bsw8pReIQgQjABDDnngivfclyxV6bUTr+vVSkSLSGWdEunpAjhCMAECM+vBDacYMV+7RI4sTXnhB+uADV37rLenIIyNaPyCnCEYAIEY9+qjbt20rHXXUAU/a+jN33OHKN94oXX55xOsH5BTBCADE6MDVYKvI449nccLdd0s7driBJM89F+nqAblCMAIAMeiff9w+KUmqXTuLJ4PdM/fem81gEiB68AkFgBgUXKG3TJksYo1rrnF7GyNyzjkRrxuQWwQjABCDpk93+3+1iuzZI02Y4MoPPUTqd8QEghEAiEFz5rh906YHPPHjj26F3oNmQQOiC8EIAMSgWbPc/vjjMxzcuzd9Bs1550kFC/pSNyC3CEYAIMYEAtLCha586qkZnvj0U2naNDeq9ZFH/KoekGsEIwAQY9askTZvdjFHnToZnhg+3O1vvllq2NCv6gG5RjACADEmOGu3Zk2X5d3zf/8nTZzoytdd51vdgLwgGAGAGDN5cuYZvJ7gDJomTaQGDXypF5BXBCMAEGN++MHtzzorw0Fbe8ZcdpkvdQLyg2AEAGLIH3+4HhlLdNa48f6DM2dK337rDmZqLgFiA8EIAMSQ9u3d/oQTpOLF9x985hm3/89/pMqVfasbkFcEIwAQI7Zvl376yZV7995/cNkyaeRIV772Wt/qBuQHwQgAxNAsGkuuWrHi/t4YSzjSs6dLAd+okdShg99VBPKEYAQAYsC2benJVW+5xeUY0fPPSx9/7A5ef72v9QPyg2AEAGJkOq8lO7MxqpbTTL/95qISc+mlUufOflcRyDOCEQCIAWPHur3FHN7A1Xbt0p984QUpJcW3ugH5RTACADHg99/T17/Tn3+6lhHz3XdS+fK+1g3IL4IRAIgBS5akp4DXU0+5ByeeKJ1yiq/1AkKBYAQAotyWLdLq1a5cM3mZ9PLL7kHfvr7WCwgVghEAiJH070ccIZUe9oi0e7fUsqV0xRV+Vw0ICYIRAIhyq1a5vbf+3aRJ7sFtt/laJyCUCEYAIMqtX+/2ZYttlxYtSl+dF4gTBCMAECPByBG7VqaPYq1Qwdc6AaFEMAIAUS44eLXsjxNcoU0bX+sDhBrBCABEud/npHr7Y/6ZJlWqJA0Y4HeVgJAiGAGAKLbm7z364Uf3p7rRYQuljz6Sypb1u1pASBGMAEC0CgT0Rss3tTtQSM0KzNAxM99m4CriEsEIAESjvXulCy7QnPmFvYeXXlVUOuYYv2sFhAXBCABEoyuvlMaP11bZqnhSmTNO8LtGQNgQjABAtPn6a+n9973i1mNO8vbeSr1AnCIYAYBo88ILbl+unDYUqeQVS5b0t0pAOBGMAEC0+f13b7ftsec1d647VK2av1UCwolgBACiydSp8iKQpCTNLnGm9uyRypeXjjvO74oB4UMwAgDR5MEH3b5rV/3415FesXlzLzYB4hbBCABEk99+c/vrrtP06a7YtKmvNQLCjmAEAKLFzz9LK1Z4xS0lq+jzz93h00/3t1pAuBGMAEC06NfP7Zs21X3DKmnzZql6dddNA8QzghEAiAZffSV9+qlXfP+KUXrmWTdI5NprpQL8pUac4yMOANHguee83bpmbdTt4epeuU4dqVcvn+sFREDBSFwEAHAQqamyASK7VVBn/zNaGzcmqVYt6ZdfpKJF/a4cEKUtI0OHDlWNGjVUpEgRNWvWTDNmzMjR60aOHKmkpCS1bds2L5cFgPhk3TPbtum5Ev00Z2Ex79CYMQQiSBy5DkZGjRql3r17a8CAAZo9e7YaNGig1q1ba82aNQd93V9//aU77rhDp512Wn7qCwDxZcsWqXt3r/hqitv37y81aOBzvYBoDkaeeuopdevWTZ07d9Zxxx2n4cOHq1ixYhoxYkS2r9m7d686dOiggQMHqpa1PQIApB07pBtvlNat0+LD6+v3deW9wzff7HfFgCgORnbt2qVZs2apVatW6W9QoID3eHowO08WHnjgAZUrV07XXXddjq6TmpqqzZs3Z9oAIO7cd5/09tte8flj3eJ4DRtKR7rEq0DCyFUwsm7dOq+Vo7wtlJCBPV61alWWr5k6dapeffVVvfzyyzm+zqBBg1SyZMm0rWrVqrmpJgBEv717pVdfdeV27fT1nhZekVYRJKKwTu3dsmWLrrnmGi8QKVu2bI5f17dvX23atCltW758eTirCQCR98470saNXnHdg8P000/u8Pnn+1stIOqn9lpAkZycrNWrV2c6bo8rVKjwr/P//PNPb+DqhRdemHZs37597sIFC2rBggWqXbv2v15XuHBhbwOAuGRL8Vo2M3P55fr8x9IKBKQTTpAqVvS7ckCUt4ykpKSocePGmjRpUqbgwh43zyJf8THHHKNff/1Vc+bMSdsuuuginXXWWV6Z7hcACenOO+VFH+bpp/Xuu67Ypo2vtQJiJ+mZTevt1KmTmjRpoqZNm2rIkCHatm2bN7vGdOzYUZUrV/bGfVgeknr16mV6falSpbz9gccBICE88og0ZIgr33671hauovHj3cMuXXytGRA7wUi7du20du1a9e/f3xu02rBhQ02cODFtUOuyZcu8GTYAgAPs2pWW9l3NmtlofS37xTWSVK4sHXWU3xUE/JEUCATbCqOXTe21WTU2mLVEiRJ+VwcA8mbUKOnKK6UjjpD+/tv6vvXJJ9JFF0n160s//+x3BQF/vr9pwgCASLDB+w8/7Mo9eniBiNmfZsRrKAESFcEIAETCM89Iv/7qyh07ph2eOtXt27XzqV5AFCAYAYBw27DBG6zq6d07bXDIpk2ut8Y0aeJj/QCfEYwAQLhZS4gNzytXzs2m2W/ePLevVEkqWdK/6gF+IxgBgHCypo9x41x50CDL6pj21A8/uP3xx/tUNyBKEIwAQDj17On2pUtL+/MxHRiMnHWWD/UCogjBCACEy5Il0pgxrvzii1JSUqant2xxe1bpRaIjGAGAcBkxwu0to9nll//r6e3b3b5YsQjXC4gyBCMAEA67d0uvvebKTzzxr1YRs3/RXh12WITrBkQZghEACEeCs+uuk1ascGNFLr74X6csXSrNnu3KtlovkMgIRgAg1Pr1k956y5WffTbLfpi5c92+dm2pVq0I1w+IMgQjABBKe/a4warmwQelq6/O8rQ1a9y+evUI1g2IUgQjABBK69e7jKumT59sT/vsM7e3PGhAoiMYAYBQ+ugjty9bVipUKNvTvvnG7Vu2jFC9gChGMAIAobJggfTkk6589tnZnrZ2rbRypSufeWaE6gZEsYJ+VwAA4sLixVKjRi55iM3VHTw421OnT08v2wBWINHRMgIAofDAAy4QOeII6fvvpYoVsz31k0/c/sYbs0w/AiQcghEAyK/586W333ZlS3RWr162p+7cKb3yiitfckmE6gdEOYIRAMivsWOlvXulZs2kCy886Klt27q95UI744zIVA+IdgQjAJAfP/4o9e3ryhdccNBTLV754gtXfuSRg062ARIKwQgA5Gf2THDWjI0R6djxoKfPmOEyxds4ka5dI1NFIBYwmwYA8uKXX6RTT5W2bHH53OfMkQ4//KAvmTo1PW4pyF9fIA2/DgCQl2m8LVpIW7e6xxMnHjIQMfPmuf0114S5fkCMoZsGAHK79oyNDbFApGZNadky6aijcvTSX391+8aNw1tFINYQjABAbvTund7EYbNoqlbN0ct273Y9O+boo8NYPyAGEYwAQE5ZkpA333Tl226T6tfP8Uv//lvatcsNXj3hhPBVEYhFBCMAkBOBgHT11dKmTVLx4tKDD+bq5cFVeuvWlQrwlxfIhF8JAMiJSZOkDz5wyUHef98FJLkQTAHfqlV4qgfEMoIRAMiJESPcvl07qXXrXL102zZp/HhXvuqqMNQNiHEEIwBwKP/7n9uCwUgu/fSTS3ZWqZLUvHnoqwfEOoIRADiUO+9MTxByiJTvWQkujNe0aYjrBcQJghEAOJhBg6QVK6SUFOnZZ/M07vWHH1yZWTRA1ghGACA7M2dK996bnl+kVKlcv8XcudL8+VKRItItt4S+ikA8IBgBgKxYUpBu3VzTho0TsWV288AWxzM2VqRs2dBWEYgXBCMAkJVHH3WL3xUuLPXr57KV5cGsWW5/0kmhrR4QTwhGAOBANnNmwABXfvppqV69PL2NzaB5/XVXZrwIkD2CEQA4cBGZhx5y5Tp1pK5d8/xW338v7djhyk2ahKh+QBwiGAGAjB54QPr9dzd75ssvXcbVPAouY3PGGdIxx4SuikC8IRgBgKCPP05vFbG1Z6pXz/NbLVggvfiiK990U4jqB8QpghEAMIsXS506ufJll0l33ZWvtwvGNKedJl1xRQjqB8Sxgn5XAACiYhpv7dqubP0pwXVo8mjdOuntt125e/cQ1A+Ic7SMAMDUqenl0aOlEiXy9XYffuj2BQtKl1ySz7oBCYBgBEBi27w5vR+lTZuQzMGdNs3tr71WKlo0328HxD2CEQCJzVpCNmxw5ZdfzvfbWW6RceNc+fzz8/12QEIgGAGQ2INWb77ZlXv0kCpWzPdbWlyzdq0rn3devt8OSAgEIwASk605YzNmUlOlBg2kxx8PyduuXu32tqaeLY4H4NAIRgAkZpbVSy+VPvjAjTK12TPFi4fkra2BxQQn5wA4NIIRAInHFowZO9aVhwyRGjUKydtOmCB9+60r9+kTkrcEEgLBCIDEsnev9Oqr6alRQ5ge9fnn3f6WW0h0BuQGwQiAxPHPP9LZZ0s//CAVKJA+eDUE5s51S9kEp/QCyDmCEQCJw6KEr76SkpOlt94K6ep1w4a5RK7HHSc1bBiytwUSAungASSGWbPcQnjGEoG0bh3SiTkTJ7py165SUlLI3hpICLSMAIh/W7e62TOmffuQBiJmxgyXssT85z8hfWsgIRCMAIhvlhL1qqukZctc4o/gKNMQByPGhqNYyhIAuUM3DYD4tWOHdMEF0uTJru/kjTekMmVCfpmFC90+RDOEgYRDMAIgfhfAa95c+v13qVAh6bXXpP/+NyyXWrDA7Y8+OixvD8S9PHXTDB06VDVq1FCRIkXUrFkzzQi2UWbh5Zdf1mmnnabSpUt7W6tWrQ56PgCEZJ7t6ae7QMSMGSN16BCWS+3Z4y5nCEaACAUjo0aNUu/evTVgwADNnj1bDRo0UOvWrbVmzZosz58yZYrat2+vr776StOnT1fVqlV1zjnnaMWKFXmsMgAcokXkjDOkn392C8TYVF7rqgmTgQOllSul0qWZ0gvkVVIgYJPScs5aQk466SQ9v38Q2L59+7wAo2fPnrr77rsP+fq9e/d6LST2+o4dO+bomps3b1bJkiW1adMmlShRIjfVBZBIbNG7Sy5xedltsOqff0qVKoXtcj/+KDVt6sqW1LVLl7BdCohJOf3+zlXLyK5duzRr1iyvqyXtDQoU8B5bq0dObN++Xbt371aZgwwiS01N9X6AjBsAHJJlVLVAJLj+TBgDkU8+kU491ZWtRYRABMi7XAUj69at81o2ypcvn+m4PV61alWO3qNPnz6qVKlSpoDmQIMGDfIiqeBmLS8AcFDr10uvvOLKTzwhtWsX1kk6N9zgMq62aJGeSw1ADOQZefTRRzVy5Eh9+OGH3uDX7PTt29dr0gluy5cvj2Q1AcQa623O2O0bwjVnstK/vxsnUqyYW/yXfy8BEZzaW7ZsWSUnJ2v16tWZjtvjChUqHPS1Tz75pBeMfPnll6pfv/5Bzy1cuLC3AUCODBokjR8vFSwoff21Gy8SxmSuQ4a4cp8+9ncxbJcCEkauWkZSUlLUuHFjTZo0Ke2YDWC1x81tPn82Hn/8cT344IOaOHGimjRpkr8aA0BGU6ZI997ryo89Jp1ySlgv17Onm85r7rgjrJcCEkauk57ZtN5OnTp5QUXTpk01ZMgQbdu2TZ07d/aetxkylStX9sZ9mMcee0z9+/fXu+++6+UmCY4tKV68uLcBQJ7Nny+ddZYrW16RXr3C3hv0/vuuPHKk66YB4EMw0q5dO61du9YLMCywaNiwodfiERzUumzZMm+GTdCwYcO8WTiXX355pvexPCX3339/CH4EAAnJIoOrr3blww6TPv007MvlWnok66axy7RtG9ZLAQkl13lG/ECeEQCZbNniVt8dN849/uCD9FV5w7jenjXCfPONdNRR6evRAMj/9zdr0wCILRYVnHmmNHu2G7D69NNhD0SMDVq1QMTcckvYLwcklIhO7QWAfOvXzwUi5tFHwz6N1yxZ4i5r7rlHuummsF8SSCgEIwBihy14ZzNmjA3auP32iFz2oYdcorO6dV2OkTAPTQESDsEIgNjw119uwKp109h4kVGjInJZmwA4YoQr2yRBUiABoUcwAiD62Th7646x5olataTXXrPER2G/7D//SBdd5Mq2Ki8zaIDwYAArgOg3eLCbOWP9I6NHR6R5Yts2yXI5LlgglSwpTZxI9wwQLgQjAKLbO+9Id97pygMHSo0bR+Sy1gtkgYj54gvppJMiclkgIdFNAyB6vfFGemKz//5Xuu++iF362Wfdvnt3AhEg3AhGAESvF190+0sucS0kEfLdd9LPP7tyxsWAAYQHwQiA6LNrl9SlizR9uktsZhnHbB+hsbL7l9qSLTB+8skRuSyQ0BgzAiC67N4t2VpWn3ziHt91l1StWsQub+NE/vjDlSdMYNAqEAm0jACIHtu3S5ddlh6I2IBVyzgWQZZh1TRsKFWqFNFLAwmLlhEA0WHvXrfGzGefuceWaSzYXxJBwfVnbOAqgMggGAHgvz17pLPPlqZMcY/ffFO65pqIV2P8eGn9ejc85aqrIn55IGHRTQPAfyNHpgcir7/uSyDy998uy7y58Ubp8MMjXgUgYRGMAPBPaqobFxIMPqybplMnX6oyebK0ebNUpoz0yCO+VAFIWHTTAPDHli1Sq1bSjBnuseVc79XLt+pYbhHToYN02GG+VQNISLSMAIg865I5/XQXiBQp4gar2vK4p57qS3WWL09fmffMM32pApDQaBkBEFlvvZWe1tRaQ778UmrSxNcq3Xaby7NmsZAlewUQWbSMAFDEUps+80z6mJAWLVyGMZ8DkY0bpTFjXLlnT5KcAX4gGAEQftYFc+WVrgnCgpJzz5U+/lgqX97vmumll9LLtIoA/iAYARBegwdLtWtLo0dLBQpId97pMqzatBWf7dsnDRrkyo89JhUq5HeNgMTEmBEA4fHVV9K997rF7kzNmtIrr0j/+Y+iwbZtLvO8ddOYm2/2u0ZA4iIYARB6zz0n3XJL+uNu3aRhw6TkZEVLwldLcBbMPG9jRYoV87tWQOKimwZA6ARnxgQDERuE8euvbmBGlAQiS5dKVaqkr8Vngcizz/pdKyCx0TICIP9sQZfevd2aMsYWd7HsYdYtY+UosWiRdMUV0urV7vGLL7pGGwD+ip6/EgBi0xdfuHTu9g1v82Jt8EX//lLZsoo2NoRlzhxXnj1bOvFEv2sEwBCMAMibvXtdIGJNDVu3Sscf71pCTj5Z0WjePOn999PX4iMQAaIHwQiA3AchL7wgPfGEy6NuGjVyi7tYavcoZVN4bSqvTebxaS0+ANkgGAGQM5Yv/eWX3WjPhQvdsVKl3LQU65aJ4kDkvfdcFnpjiwQDiC4EIwAO7q+/pCefdEnL1q51x2xQqmUJu/HGqA5CzNSpLl4yTZv6thYfgIMgGAGQvWXLpNNOk/7v/9zjcuWkfv2kq6+WSpdWtLNBqhdc4HqWWraUPvrI7xoByArBCICsjRwpXX+9tGWLdMQRrp/DvtFTUhQLfvpJuugiadMmqVYt9+McdpjftQKQFZKeAcg6Z4j1bVggcvTRbnDqeefFRCCyc6cb1mKTelaskCpUcF01UTjTGMB+tIwASB+gauNAbNrJjh3umK20a/NgCxdWLLDYqVUracYM97huXWnKFBeQAIhetIwAic6aEp5+WjruODcrxgKRY45xa8m8+27MBCI2trZBg/RAZMAAaeZMAhEgFtAyAiSq1FSXpMxyhvz+uztmgUevXm7+awx0yQTNn++6ZWx8iBk3Tjr/fL9rBSCnCEaARGOpSG3hOktHGpwlU7Kk9OCDLq275Q6JIdu2SZdd5gIRm+BjM5CtqwZA7CAYARLB7t3SN9+47hhrNggqU0a66y63WpyVY5Algg027Pz4o1S7tt81ApBbBCNAPFuwQBoxQnrttfSEZcZyh1x3nXTxxTHXEhIUCEi33+7iK9OjB4EIEKsIRoB4Y9/SX33lZsZ8/nn6ceuKsdkxljvE1pKJcZb8dfhwV7a1Zh5/3O8aAcgrghEgHmzfLn3wgfTDD9KkSW5EZ5CtDGc5Qm66SSpaVPHAftRgIGIJYW24C4DYRTACxPq6MS++6Baws2RlQbZ2jI0Due02l7QsznqeunRx5Xr1pPvv97tGAPKLYASINbbQii2yYtNvf/kl/XilSq4bpnFj1xoSZwk2rPfJWkR69pQ2b3apUGzAanKy3zUDkF8EI0As5QV54w2XIdVaRIIswYZNybXmgihfQTevbAaydcfYjx80dGjc/rhAwiEYAaKZpRAdP1768ktp2jTXKmJsBswNN7hmgsqVFc9sCEzbttLWrVKBAtItt0gdOkhNmvhdMwChQjACRGMzwMcfS//7n1vhLSPLBWLjQGxOa7Fiivcs9dYaMniwe1ynjstQT0IzIP4QjAB+27fPjf34+mvpk09cU0BQUpLLBXLOOdJZZ7lvZBucGud++sn1PP32m3t8ySUuaSwr7wLxKf7/qgHRaM0a6bvvXMuHtYCsXJn5+YYNpTZtpGuvdQFIgrBMqs88I736quuRKlfOLZ9z4YV+1wxAOBGMAJEYeDpnjjR7tltS1oKQP/7IfI6NxDz9dNf6cemlcTcd91ANQ9Yg9Oij0vffpx+322C5RI480s/aAYgEghEg1JYvl6ZMcYNPLfiwPgcLSA50/PFSixYuCLFv3jhJSJZTu3ZJX3whPfWUNHmyO2Y9UBdcIN18s9Sypd81BBApBCNAfjOfWquHDW6wwMPSsP/557/Ps+VkbQruiSe6AKR5c3csAdkSOdYNYwNT161zxwoXlq6+WnrgAZcuBUBiIRgBDtWHYOM5li6Vli1z+yVLpHnzXCrQ1av//Rqbf2rzTi3gsH2zZm7chw1GTUCWrMxulcVpb7/tZigHWV62K65wmerr1vWzlgD8RDAC2EhJG8OxaJFr1Vi82O1ts8Ajqy6WjOwbtX59qUED1+ViK+LaonQJats26ddf3eQgWypn+nRp1arM51jgYYsG9+qVEJODABwCfwaQOP88txzi1sphXSoWfCxc6Fo75s7NuoUjyPKNV6kiVa8uVasm1ajhvk0tH3mtWq67JUFbPSzIsF6qjJvdVrvdGdntsd4pyxHy3/9Kxx7rV40BRCOCEcQ+a7lYscJtf/+dvlnysOBmj23EZHYOO0w66iipdm23WZARLFetyj/fJe3Y4RqKrMHIxueOHu1ubXaNRdZDdeaZUqNGrqcqznO0AciHPP2FHTp0qJ544gmtWrVKDRo00HPPPaemTZtme/57772n++67T3/99ZeOOuooPfbYYzr//PPzU2/E+zgNa8X45x83wtFWo7Vc4Jabw1oybASkla2VwzYLNHKqeHH3z3Jr2QgGG9bSYd+WKSlKdNZjZQGG9VRZr5UNjbFt/nw3XCarFg+7lZYWJbhZb1WcrdEHINqCkVGjRql3794aPny4mjVrpiFDhqh169ZasGCBylmGogNMmzZN7du316BBg3TBBRfo3XffVdu2bTV79mzVs/W/ET8BhP3T2TabYRIsBzcLJizAsM3KwW3TJmnjRmnDBtfmb90ldszeLzcsT4et0WJTMYJ761qxsm1Wts9nAkyftYDB/hfYbcxqs9u9ZYu7/Va2Wx7c7H/Bnj3Zv3eJEq7RyNKgXHSRWzPGGpUAID+SAoED/61zcBaAnHTSSXr++ee9x/v27VPVqlXVs2dP3X333f86v127dtq2bZs+/fTTtGMnn3yyGjZs6AU0ObF582aVLFlSmzZtUgn7a4h09qVt/5y1b5Dgtnu3W9jDui9sC5bzcswCCRuRmF2QETx2sC6QvLIA44gjXA5w+/9u67LYuI3y5d2xihXdP8GtZcMeR/m4DftNs/89dlvtdgVvcW7LdruDcV0wwAiWLaazoCI//zsKFZJq1nRBhw2LsYYk26wcA7cZQBTJ6fd3rlpGdu3apVmzZqlv375pxwoUKKBWrVppug2Zz4Idt5aUjKwlZezYsdleJzU11dsy/jDhMOTir7w+cNf2HFDA+8d4wD32Ntu5ffCYOzXDOcp4TLk4J/P7ese8y2dxvbQ6SYF9AReAWGX37X9vZf52OPBxVsfSHxfevymf75PhWIFkN+gzuaBUMFkBb19QKpSyf19QgYLpZeseCaQUlgoXcQGIdZfYMXufjO9rP67NysgwMyOrUPrAYzk9x25rxs1ivKzKOXls8WBWAUXuQv/8sRnG9rtvE3sybrbgr+2tx8qet9jOGo0srrPNGpXsfx8AREqugpF169Zp7969Km9/vTKwx/OtUzkLNq4kq/PteHasS2fgwIEKt9GTy2r61hPCfp2Es2//ttvvikQ3CxYs2ZfFXrbPrnzgMetpOvzw9OAiY8BhE3vs180akSzYoBUDQCyIyikC1vKSsTXFWkasKyjUOrXdpP8s/cb9xd6/eX+8Mzz2jhXIfCzpgOfdOQd/TVavy/SaAlmck9Wx5AJKsm8x25Jtn+yVk5KDjwvILu699wEO/GLK6osqHs7J6+u8W+puZ9p2sMcHe84afQ4VVDBBBwCcXP05LFu2rJKTk7X6gJwM9rhCNsPn7XhuzjeFCxf2tnC74a1Tw34NAABwcPbv9xxLSUlR48aNNWnSpLRjNoDVHje31NdZsOMZzzdffPFFtucDAIDEkuuGYus+6dSpk5o0aeLlFrGpvTZbpnPnzt7zHTt2VOXKlb1xH+bWW2/VGWecocGDB6tNmzYaOXKkZs6cqZdeein0Pw0AAIj/YMSm6q5du1b9+/f3BqHaFN2JEyemDVJdtmyZN8Mm6JRTTvFyi/Tr10/33HOPl/TMZtKQYwQAAOQpz4gfyDMCAEDsyen3d67GjAAAAIQawQgAAPAVwQgAAPAVwQgAAPAVwQgAAPAVwQgAAPAVwQgAAPAVwQgAAPAVwQgAAPBVTCxiHkwSa5ncAABAbAh+bx8q2XtMBCNbtmzx9lWrVvW7KgAAIA/f45YWPqbXptm3b5/+/vtvHX744UpKSgppxGYBzvLly1nz5hC4V7nD/co57lXOca9yjnsVHffKQgwLRCpVqpRpEd2YbBmxH6BKlSphe3+7+XxYc4Z7lTvcr5zjXuUc9yrnuFf+36uDtYgEMYAVAAD4imAEAAD4KqGDkcKFC2vAgAHeHgfHvcod7lfOca9yjnuVc9yr2LpXMTGAFQAAxK+EbhkBAAD+IxgBAAC+IhgBAAC+IhgBAAC+SuhgZOjQoapRo4aKFCmiZs2aacaMGUok999/v5fRNuN2zDHHpD2/c+dO3XTTTTriiCNUvHhxXXbZZVq9enWm91i2bJnatGmjYsWKqVy5crrzzju1Z88exYNvvvlGF154oZc50O7N2LFjMz1vY7/79++vihUrqmjRomrVqpX++OOPTOds2LBBHTp08BIJlSpVStddd522bt2a6ZxffvlFp512mvc5tCyIjz/+uOLtXl177bX/+qyde+65CXevBg0apJNOOsnLJm2/L23bttWCBQsynROq37spU6aoUaNG3gyJOnXq6PXXX1esycn9OvPMM//12erevXvC3a9hw4apfv36aYnLmjdvrgkTJsTO5yqQoEaOHBlISUkJjBgxIvDbb78FunXrFihVqlRg9erVgUQxYMCAwPHHHx9YuXJl2rZ27dq057t37x6oWrVqYNKkSYGZM2cGTj755MApp5yS9vyePXsC9erVC7Rq1Srw008/BcaPHx8oW7ZsoG/fvoF4YD/PvffeGxgzZozNOAt8+OGHmZ5/9NFHAyVLlgyMHTs28PPPPwcuuuiiQM2aNQM7duxIO+fcc88NNGjQIPD9998Hvv3220CdOnUC7du3T3t+06ZNgfLlywc6dOgQmDt3buB///tfoGjRooEXX3wxEE/3qlOnTt69yPhZ27BhQ6ZzEuFetW7dOvDaa6959Z8zZ07g/PPPD1SrVi2wdevWkP7eLV68OFCsWLFA7969A7///nvgueeeCyQnJwcmTpwYiCU5uV9nnHGG9/c742fLPiuJdr8+/vjjwLhx4wILFy4MLFiwIHDPPfcEChUq5N27WPhcJWww0rRp08BNN92U9njv3r2BSpUqBQYNGhRIpGDE/vhnZePGjd4H+b333ks7Nm/ePO+LZvr06d5j+7AWKFAgsGrVqrRzhg0bFihRokQgNTU1EE8O/ILdt29foEKFCoEnnngi0z0rXLiw9yVp7JfVXvfjjz+mnTNhwoRAUlJSYMWKFd7jF154IVC6dOlM96tPnz6BunXrBmJVdsHIxRdfnO1rEvVerVmzxvu5v/7665D+3t11113ePzQyateunfflHssOvF/BYOTWW2/N9jWJfL9Kly4deOWVV2Lic5WQ3TS7du3SrFmzvGb1jOvf2OPp06crkVi3gjWt16pVy2sit2Y6Y/dn9+7dme6RdeFUq1Yt7R7Z/oQTTlD58uXTzmndurW36NJvv/2meLZkyRKtWrUq0/2x9Resuy/j/bHuhiZNmqSdY+fbZ+2HH35IO+f0009XSkpKpntoTdH//POP4ok171rTb926ddWjRw+tX78+7blEvVebNm3y9mXKlAnp752dk/E9gufE+t+3A+9X0DvvvKOyZcuqXr166tu3r7Zv3572XCLer71792rkyJHatm2b110TC5+rmFgoL9TWrVvn/c/KeNONPZ4/f74ShX1xWn+ffTmsXLlSAwcO9Prj586d633R2h99+4I48B7Zc8b2Wd3D4HPxLPjzZfXzZ7w/9uWbUcGCBb0/pBnPqVmz5r/eI/hc6dKlFQ9sfMill17q/ax//vmn7rnnHp133nneH7Hk5OSEvFe2Gvltt92mFi1aeF+iJlS/d9mdY18sO3bs8MY4xZqs7pe56qqrVL16de8fVTamqE+fPl6AOmbMmIS7X7/++qsXfNj4EBsX8uGHH+q4447TnDlzov5zlZDBCBz7MgiygU8WnNgv9ejRo2Pmlw+x4corr0wr27++7PNWu3Ztr7WkZcuWSkQ2mNAC/6lTp/pdlZi+X9dff32mz5YNKLfPlAW99hlLJHXr1vUCD2tBev/999WpUyd9/fXXigUJ2U1jzXn2r7EDRxLb4woVKihRWdR89NFHa9GiRd59sO6sjRs3ZnuPbJ/VPQw+F8+CP9/BPkO2X7NmTabnbWS6zRpJ9Hto3YL2e2iftUS8VzfffLM+/fRTffXVV6pSpUra8VD93mV3js2yiMV/aGR3v7Ji/6gyGT9biXK/UlJSvBkujRs39mYiNWjQQM8880xMfK4SMhix/2H2P2vSpEmZmgDtsTVxJSqbRmn/mrB/Wdj9KVSoUKZ7ZE2fNqYkeI9sb82CGb9EvvjiC++DaU2D8cy6C+wXM+P9saZKG9+Q8f7YL7/11wZNnjzZ+6wF/2DaOTYt1vpzM95D+xdOrHU75Mb//d//eWNG7LOWSPfKxvfaF6s1n9vPd2C3U6h+7+ycjO8RPCfW/r4d6n5lxVoGTMbPVqLcrwPZ709qampsfK4CCTy112Y+vP76695I/uuvv96b2ptxJHG8u/322wNTpkwJLFmyJPDdd995U7psKpeNWA9OBbNpdJMnT/amgjVv3tzbDpwKds4553jT7mx615FHHhk3U3u3bNniTXGzzX5VnnrqKa+8dOnStKm99pn56KOPAr/88os3WySrqb0nnnhi4IcffghMnTo1cNRRR2Warmqj3G266jXXXONNwbPPpU2di6Xpqoe6V/bcHXfc4Y3at8/al19+GWjUqJF3L3bu3JlQ96pHjx7edHD7vcs4FXX79u1p54Ti9y44BfPOO+/0Zk0MHTo05qaq5uR+LVq0KPDAAw9498k+W/a7WKtWrcDpp5+ecPfr7rvv9mYZ2X2wv0f22Gajff755zHxuUrYYMTYHGn7n2P5Rmyqr+U3SCQ2JatixYrez1+5cmXvsf1yB9mX6o033uhND7MP4CWXXOL9Icjor7/+Cpx33nlevgcLZCzA2b17dyAefPXVV94X64GbTVMNTu+97777vC9IC2xbtmzpze/PaP369d4XavHixb0pcp07d/a+nDOyHCWnnnqq9x72/8GCnHi6V/bFYX/g7A+bTS+sXr26lxfiwMA/Ee5VVvfINsulEerfO/t/0rBhQ+/3276gM14jXu7XsmXLvMCjTJky3mfCctPYF2XGPCOJcr+6dOni/W5Z/e13zf4eBQORWPhcJdl/8t++AgAAkDcJOWYEAABED4IRAADgK4IRAADgK4IRAADgK4IRAADgK4IRAADgK4IRAADgK4IRAADgK4IRAADgK4IRAADgK4IRAADgK4IRAAAgP/0/JAxY/9RBujIAAAAASUVORK5CYII=",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "n_samples = 1000\n",
    "n_classes = 3\n",
    "\n",
    "t = (torch.rand(n_samples, n_classes) - .5) * 10\n",
    "probas = F.softmax(t, -1)\n",
    "sharpened_probas = Sharpen()(probas)\n",
    "plt.plot(probas.flatten().sort().values, color='r')\n",
    "plt.plot(sharpened_probas.flatten().sort().values, color='b')\n",
    "plt.show()\n",
    "test_gt(sharpened_probas[n_samples//2:].max(-1).values.sum().item(), probas[n_samples//2:].max(-1).values.sum().item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class Sequential(nn.Sequential):\n",
    "    \"\"\"Class that allows you to pass one or multiple inputs\"\"\"\n",
    "    def forward(self, *x):\n",
    "        for i, module in enumerate(self._modules.values()):\n",
    "            x = module(*x) if isinstance(x, (list, tuple, L)) else module(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class TimeDistributed(nn.Module):\n",
    "    def __init__(self, module, batch_first=False):\n",
    "        super(TimeDistributed, self).__init__()\n",
    "        self.module = module\n",
    "        self.batch_first = batch_first\n",
    "\n",
    "    def forward(self, x):\n",
    "\n",
    "        if len(x.size()) <= 2:\n",
    "            return self.module(x)\n",
    "\n",
    "        # Squash samples and timesteps into a single axis\n",
    "        x_reshape = x.contiguous().reshape(-1, x.size(-1))  # (samples * timesteps, input_size)\n",
    "\n",
    "        y = self.module(x_reshape)\n",
    "\n",
    "        # We have to reshape Y\n",
    "        if self.batch_first:\n",
    "            y = y.contiguous().reshape(x.size(0), -1, y.size(-1))  # (samples, timesteps, output_size)\n",
    "        else:\n",
    "            y = y.reshape(-1, x.size(1), y.size(-1))  # (timesteps, samples, output_size)\n",
    "\n",
    "        return y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class Temp_Scale(Module):\n",
    "    \"Used to perform Temperature Scaling (dirichlet=False) or Single-parameter Dirichlet calibration (dirichlet=True)\"\n",
    "    def __init__(self, temp=1., dirichlet=False):\n",
    "        self.weight = nn.Parameter(tensor(temp))\n",
    "        self.bias = None\n",
    "        self.log_softmax = dirichlet\n",
    "\n",
    "    def forward(self, x):\n",
    "        if self.log_softmax: x = F.log_softmax(x, dim=-1)\n",
    "        return x.div(self.weight)\n",
    "\n",
    "\n",
    "class Vector_Scale(Module):\n",
    "    \"Used to perform Vector Scaling (dirichlet=False) or Diagonal Dirichlet calibration (dirichlet=True)\"\n",
    "    def __init__(self, n_classes=1, dirichlet=False):\n",
    "        self.weight = nn.Parameter(torch.ones(n_classes))\n",
    "        self.bias = nn.Parameter(torch.zeros(n_classes))\n",
    "        self.log_softmax = dirichlet\n",
    "\n",
    "    def forward(self, x):\n",
    "        if self.log_softmax: x = F.log_softmax(x, dim=-1)\n",
    "        return x.mul(self.weight).add(self.bias)\n",
    "\n",
    "\n",
    "class Matrix_Scale(Module):\n",
    "    \"Used to perform Matrix Scaling (dirichlet=False) or Dirichlet calibration (dirichlet=True)\"\n",
    "    def __init__(self, n_classes=1, dirichlet=False):\n",
    "        self.ms = nn.Linear(n_classes, n_classes)\n",
    "        self.ms.weight.data = nn.Parameter(torch.eye(n_classes))\n",
    "        nn.init.constant_(self.ms.bias.data, 0.)\n",
    "        self.weight = self.ms.weight\n",
    "        self.bias = self.ms.bias\n",
    "        self.log_softmax = dirichlet\n",
    "\n",
    "    def forward(self, x):\n",
    "        if self.log_softmax: x = F.log_softmax(x, dim=-1)\n",
    "        return self.ms(x)\n",
    "\n",
    "\n",
    "def get_calibrator(calibrator=None, n_classes=1, **kwargs):\n",
    "    if calibrator is None or not calibrator: return noop\n",
    "    elif calibrator.lower() == 'temp': return Temp_Scale(dirichlet=False, **kwargs)\n",
    "    elif calibrator.lower() == 'vector': return Vector_Scale(n_classes=n_classes, dirichlet=False, **kwargs)\n",
    "    elif calibrator.lower() == 'matrix': return Matrix_Scale(n_classes=n_classes, dirichlet=False, **kwargs)\n",
    "    elif calibrator.lower() == 'dtemp': return Temp_Scale(dirichlet=True, **kwargs)\n",
    "    elif calibrator.lower() == 'dvector': return Vector_Scale(n_classes=n_classes, dirichlet=True, **kwargs)\n",
    "    elif calibrator.lower() == 'dmatrix': return Matrix_Scale(n_classes=n_classes, dirichlet=True, **kwargs)\n",
    "    else: assert False, f'please, select a correct calibrator instead of {calibrator}'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bs = 2\n",
    "c_out = 3\n",
    "\n",
    "t = torch.rand(bs, c_out)\n",
    "for calibrator, cal_name in zip(['temp', 'vector', 'matrix'], ['Temp_Scale', 'Vector_Scale', 'Matrix_Scale']):\n",
    "    cal = get_calibrator(calibrator, n_classes=c_out)\n",
    "#     print(calibrator)\n",
    "#     print(cal.weight, cal.bias, '\\n')\n",
    "    test_eq(cal(t), t)\n",
    "    test_eq(cal.__class__.__name__, cal_name)\n",
    "for calibrator, cal_name in zip(['dtemp', 'dvector', 'dmatrix'], ['Temp_Scale', 'Vector_Scale', 'Matrix_Scale']):\n",
    "    cal = get_calibrator(calibrator, n_classes=c_out)\n",
    "#     print(calibrator)\n",
    "#     print(cal.weight, cal.bias, '\\n')\n",
    "    test_eq(cal(t), F.log_softmax(t, dim=1))\n",
    "    test_eq(cal.__class__.__name__, cal_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bs = 2\n",
    "c_out = 3\n",
    "\n",
    "t = torch.rand(bs, c_out)\n",
    "\n",
    "test_eq(Temp_Scale()(t).shape, t.shape)\n",
    "test_eq(Vector_Scale(c_out)(t).shape, t.shape)\n",
    "test_eq(Matrix_Scale(c_out)(t).shape, t.shape)\n",
    "test_eq(Temp_Scale(dirichlet=True)(t).shape, t.shape)\n",
    "test_eq(Vector_Scale(c_out, dirichlet=True)(t).shape, t.shape)\n",
    "test_eq(Matrix_Scale(c_out, dirichlet=True)(t).shape, t.shape)\n",
    "\n",
    "test_eq(Temp_Scale()(t), t)\n",
    "test_eq(Vector_Scale(c_out)(t), t)\n",
    "test_eq(Matrix_Scale(c_out)(t), t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bs = 2\n",
    "c_out = 5\n",
    "\n",
    "t = torch.rand(bs, c_out)\n",
    "test_eq(Vector_Scale(c_out)(t), t)\n",
    "test_eq(Vector_Scale(c_out).weight.data, torch.ones(c_out))\n",
    "test_eq(Vector_Scale(c_out).weight.requires_grad, True)\n",
    "test_eq(type(Vector_Scale(c_out).weight), torch.nn.parameter.Parameter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bs = 2\n",
    "c_out = 3\n",
    "weight = 2\n",
    "bias = 1\n",
    "\n",
    "t = torch.rand(bs, c_out)\n",
    "test_eq(Matrix_Scale(c_out)(t).shape, t.shape)\n",
    "test_eq(Matrix_Scale(c_out).weight.requires_grad, True)\n",
    "test_eq(type(Matrix_Scale(c_out).weight), torch.nn.parameter.Parameter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class LogitAdjustmentLayer(Module):\n",
    "    \"Logit Adjustment for imbalanced datasets\"\n",
    "    def __init__(self, class_priors):\n",
    "        self.class_priors = class_priors\n",
    "    def forward(self, x):\n",
    "        return x.add(self.class_priors)\n",
    "\n",
    "LogitAdjLayer = LogitAdjustmentLayer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bs, n_classes = 16, 3\n",
    "class_priors = torch.rand(n_classes)\n",
    "logits = torch.randn(bs, n_classes) * 2\n",
    "test_eq(LogitAdjLayer(class_priors)(logits), logits + class_priors)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class PPV(Module):\n",
    "    def __init__(self, dim=-1):\n",
    "        self.dim = dim\n",
    "    def forward(self, x):\n",
    "        return torch.gt(x, 0).sum(dim=self.dim).float() / x.shape[self.dim]\n",
    "    def __repr__(self): return f'{self.__class__.__name__}(dim={self.dim})'\n",
    "\n",
    "\n",
    "class PPAuc(Module):\n",
    "    def __init__(self, dim=-1):\n",
    "        self.dim = dim\n",
    "    def forward(self, x):\n",
    "        x = F.relu(x).sum(self.dim) / (abs(x).sum(self.dim) + 1e-8)\n",
    "        return x\n",
    "    def __repr__(self): return f'{self.__class__.__name__}(dim={self.dim})'\n",
    "\n",
    "\n",
    "class MaxPPVPool1d(Module):\n",
    "    \"Drop-in replacement for AdaptiveConcatPool1d - multiplies nf by 2\"\n",
    "    def forward(self, x):\n",
    "        _max = x.max(dim=-1).values\n",
    "        _ppv = torch.gt(x, 0).sum(dim=-1).float() / x.shape[-1]\n",
    "        return torch.cat((_max, _ppv), dim=-1).unsqueeze(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bs = 2\n",
    "nf = 5\n",
    "sl = 4\n",
    "\n",
    "t = torch.rand(bs, nf, sl)\n",
    "test_eq(MaxPPVPool1d()(t).shape, (bs, nf*2, 1))\n",
    "test_eq(MaxPPVPool1d()(t).shape, AdaptiveConcatPool1d(1)(t).shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class AdaptiveWeightedAvgPool1d(Module):\n",
    "    '''Global Pooling layer that performs a weighted average along the temporal axis\n",
    "\n",
    "    It can be considered as a channel-wise form of local temporal attention. Inspired by the paper:\n",
    "    Hyun, J., Seong, H., & Kim, E. (2019). Universal Pooling--A New Pooling Method for Convolutional Neural Networks. arXiv preprint arXiv:1907.11440.'''\n",
    "\n",
    "    def __init__(self, n_in, seq_len, mult=2, n_layers=2, ln=False, dropout=0.5, act=nn.ReLU(), zero_init=True):\n",
    "        layers = nn.ModuleList()\n",
    "        for i in range(n_layers):\n",
    "            inp_mult = mult if i > 0 else 1\n",
    "            out_mult = mult if i < n_layers -1 else 1\n",
    "            p = dropout[i] if is_listy(dropout) else dropout\n",
    "            layers.append(LinLnDrop(seq_len * inp_mult, seq_len * out_mult, ln=False, p=p,\n",
    "                                    act=act if i < n_layers-1 and n_layers > 1 else None))\n",
    "        self.layers = layers\n",
    "        self.softmax = SoftMax(-1)\n",
    "        if zero_init: init_lin_zero(self)\n",
    "\n",
    "    def forward(self, x):\n",
    "        wap = x\n",
    "        for l in self.layers: wap = l(wap)\n",
    "        wap = self.softmax(wap)\n",
    "        return torch.mul(x, wap).sum(-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class GAP1d(Module):\n",
    "    \"Global Adaptive Pooling + Flatten\"\n",
    "    def __init__(self, output_size=1):\n",
    "        self.gap = nn.AdaptiveAvgPool1d(output_size)\n",
    "        self.flatten = Reshape()\n",
    "    def forward(self, x):\n",
    "        return self.flatten(self.gap(x))\n",
    "\n",
    "\n",
    "class GACP1d(Module):\n",
    "    \"Global AdaptiveConcatPool + Flatten\"\n",
    "    def __init__(self, output_size=1):\n",
    "        self.gacp = AdaptiveConcatPool1d(output_size)\n",
    "        self.flatten = Reshape()\n",
    "    def forward(self, x):\n",
    "        return self.flatten(self.gacp(x))\n",
    "\n",
    "\n",
    "class GAWP1d(Module):\n",
    "    \"Global AdaptiveWeightedAvgPool1d + Flatten\"\n",
    "    def __init__(self, n_in, seq_len, n_layers=2, ln=False, dropout=0.5, act=nn.ReLU(), zero_init=False):\n",
    "        self.gacp = AdaptiveWeightedAvgPool1d(n_in, seq_len, n_layers=n_layers, ln=ln, dropout=dropout, act=act, zero_init=zero_init)\n",
    "        self.flatten = Reshape()\n",
    "    def forward(self, x):\n",
    "        return self.flatten(self.gacp(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class GlobalWeightedAveragePool1d(Module):\n",
    "    \"\"\" Global Weighted Average Pooling layer\n",
    "\n",
    "    Inspired by Building Efficient CNN Architecture for Offline Handwritten Chinese Character Recognition\n",
    "    https://arxiv.org/pdf/1804.01259.pdf\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, n_in, seq_len):\n",
    "        self.weight = nn.Parameter(torch.ones(1, n_in, seq_len))\n",
    "        self.bias = nn.Parameter(torch.zeros(1, n_in, seq_len))\n",
    "\n",
    "    def forward(self, x):\n",
    "        α = F.softmax(torch.sigmoid(x * self.weight + self.bias), dim=-1)\n",
    "        return (x * α).sum(-1)\n",
    "\n",
    "GWAP1d = GlobalWeightedAveragePool1d\n",
    "\n",
    "def gwa_pool_head(n_in, c_out, seq_len, bn=True, fc_dropout=0.):\n",
    "    return nn.Sequential(GlobalWeightedAveragePool1d(n_in, seq_len), Reshape(), LinBnDrop(n_in, c_out, p=fc_dropout, bn=bn))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = torch.randn(16, 64, 50)\n",
    "head = gwa_pool_head(64, 5, 50)\n",
    "test_eq(head(t).shape, (16, 5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class AttentionalPool1d(Module):\n",
    "    \"\"\"Global Adaptive Pooling layer inspired by Attentional Pooling for Action Recognition https://arxiv.org/abs/1711.01467\"\"\"\n",
    "    def __init__(self, n_in, c_out, bn=False):\n",
    "        store_attr()\n",
    "        self.bn = nn.BatchNorm1d(n_in) if bn else None\n",
    "        self.conv1 = Conv1d(n_in, 1, 1)\n",
    "        self.conv2 = Conv1d(n_in, c_out, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        if self.bn is not None: x = self.bn(x)\n",
    "        return (self.conv1(x) @ self.conv2(x).transpose(1,2)).transpose(1,2)\n",
    "\n",
    "class GAttP1d(nn.Sequential):\n",
    "    def __init__(self, n_in, c_out, bn=False):\n",
    "        super().__init__(AttentionalPool1d(n_in, c_out, bn=bn), Reshape())\n",
    "\n",
    "def attentional_pool_head(n_in, c_out, seq_len=None, bn=True, **kwargs):\n",
    "    return nn.Sequential(AttentionalPool1d(n_in, c_out, bn=bn, **kwargs), Reshape())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bs, c_in, seq_len = 16, 1, 50\n",
    "c_out = 3\n",
    "t = torch.rand(bs, c_in, seq_len)\n",
    "test_eq(GAP1d()(t).shape, (bs, c_in))\n",
    "test_eq(GACP1d()(t).shape, (bs, c_in*2))\n",
    "bs, c_in, seq_len = 16, 4, 50\n",
    "t = torch.rand(bs, c_in, seq_len)\n",
    "test_eq(GAP1d()(t).shape, (bs, c_in))\n",
    "test_eq(GACP1d()(t).shape, (bs, c_in*2))\n",
    "test_eq(GAWP1d(c_in, seq_len, n_layers=2, ln=False, dropout=0.5, act=nn.ReLU(), zero_init=False)(t).shape, (bs, c_in))\n",
    "test_eq(GAWP1d(c_in, seq_len, n_layers=2, ln=False, dropout=0.5, act=nn.ReLU(), zero_init=False)(t).shape, (bs, c_in))\n",
    "test_eq(GAWP1d(c_in, seq_len, n_layers=1, ln=False, dropout=0.5, zero_init=False)(t).shape, (bs, c_in))\n",
    "test_eq(GAWP1d(c_in, seq_len, n_layers=1, ln=False, dropout=0.5, zero_init=True)(t).shape, (bs, c_in))\n",
    "test_eq(AttentionalPool1d(c_in, c_out)(t).shape, (bs, c_out, 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bs, c_in, seq_len = 16, 128, 50\n",
    "c_out = 14\n",
    "t = torch.rand(bs, c_in, seq_len)\n",
    "attp = attentional_pool_head(c_in, c_out)\n",
    "test_eq(attp(t).shape, (bs, c_out))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class PoolingLayer(Module):\n",
    "    def __init__(self, method='cls', seq_len=None, token=True, seq_last=True):\n",
    "        method = method.lower()\n",
    "        assert method in ['cls', 'max', 'mean', 'max-mean', 'linear', 'conv1d', 'flatten']\n",
    "        if method == 'cls': assert token, 'you can only choose method=cls if a token exists'\n",
    "        self.method = method\n",
    "        self.token = token\n",
    "        self.seq_last = seq_last\n",
    "        if method == 'linear' or method == 'conv1d':\n",
    "            self.linear = nn.Linear(seq_len - token, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        if self.method == 'cls':\n",
    "            return x[..., 0] if self.seq_last else x[:, 0]\n",
    "        if self.token:\n",
    "            x = x[..., 1:] if self.seq_last else x[:, 1:]\n",
    "        if self.method == 'max':\n",
    "            return torch.max(x, -1)[0] if self.seq_last else torch.max(x, 1)[0]\n",
    "        elif self.method == 'mean':\n",
    "            return torch.mean(x, -1) if self.seq_last else torch.mean(x, 1)\n",
    "        elif self.method == 'max-mean':\n",
    "            return torch.cat([torch.max(x, -1)[0] if self.seq_last else torch.max(x, 1)[0],\n",
    "                              torch.mean(x, -1) if self.seq_last else torch.mean(x, 1)], 1)\n",
    "        elif self.method == 'flatten':\n",
    "            return x.flatten(1)\n",
    "        elif self.method == 'linear' or self.method == 'conv1d':\n",
    "            return self.linear(x)[...,0] if self.seq_last else self.linear(x.transpose(1,2))[...,0]\n",
    "\n",
    "    def __repr__(self): return f\"{self.__class__.__name__}(method={self.method}, token={self.token}, seq_last={self.seq_last})\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = torch.arange(24).reshape(2, 3, 4).float()\n",
    "test_eq(PoolingLayer('cls', token=True, seq_last=True)(t), tensor([[ 0.,  4.,  8.], [12., 16., 20.]]))\n",
    "test_eq(PoolingLayer('max', token=True, seq_last=True)(t), tensor([[ 3.,  7., 11.], [15., 19., 23.]]))\n",
    "test_close(PoolingLayer('mean', token=True, seq_last=True)(t), tensor([[ 2.,  6., 10.], [14., 18., 22.]]))\n",
    "test_close(PoolingLayer('max-mean', token=True, seq_last=True)(t), tensor([[ 3.,  7., 11.,  2.,  6., 10.],\n",
    "                                                                           [15., 19., 23., 14., 18., 22.]]))\n",
    "test_close(PoolingLayer('flatten', token=True, seq_last=True)(t), tensor([[ 1.,  2.,  3.,  5.,  6.,  7.,  9., 10., 11.],\n",
    "                                                                          [13., 14., 15., 17., 18., 19., 21., 22., 23.]]))\n",
    "test_eq(PoolingLayer('linear', seq_len=4, token=True, seq_last=True)(t).shape, (2, 3))\n",
    "test_eq(PoolingLayer('max', token=False, seq_last=True)(t), tensor([[ 3.,  7., 11.], [15., 19., 23.]]))\n",
    "test_close(PoolingLayer('mean', token=False, seq_last=True)(t), tensor([[ 1.5000,  5.5000,  9.5000],\n",
    "                                                                        [13.5000, 17.5000, 21.5000]]))\n",
    "test_close(PoolingLayer('max-mean', token=False, seq_last=True)(t), tensor([[ 3.,  7., 11.,  1.5000,  5.5000,  9.5000],\n",
    "                                                                            [15., 19., 23., 13.5000, 17.5000, 21.5000]]))\n",
    "test_close(PoolingLayer('flatten', token=False, seq_last=True)(t), tensor([[ 0.,  1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 10., 11.],\n",
    "                                                                           [12., 13., 14., 15., 16., 17., 18., 19., 20., 21., 22., 23.]]))\n",
    "test_eq(PoolingLayer('linear', seq_len=4, token=False, seq_last=True)(t).shape, (2, 3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = torch.arange(24).reshape(2, 3, 4).swapaxes(1,2).float()\n",
    "test_eq(PoolingLayer('cls', token=True, seq_last=False)(t), tensor([[ 0.,  4.,  8.], [12., 16., 20.]]))\n",
    "test_eq(PoolingLayer('max', token=True, seq_last=False)(t), tensor([[ 3.,  7., 11.], [15., 19., 23.]]))\n",
    "test_close(PoolingLayer('mean', token=True, seq_last=False)(t), tensor([[ 2.,  6., 10.], [14., 18., 22.]]))\n",
    "test_close(PoolingLayer('max-mean', token=True, seq_last=False)(t), tensor([[ 3.,  7., 11.,  2.,  6., 10.],\n",
    "                                                                           [15., 19., 23., 14., 18., 22.]]))\n",
    "test_close(PoolingLayer('flatten', token=True, seq_last=False)(t), tensor([[ 1.,  5.,  9.,  2.,  6., 10.,  3.,  7., 11.],\n",
    "                                                                           [13., 17., 21., 14., 18., 22., 15., 19., 23.]]))\n",
    "t = torch.arange(24).reshape(2, 3, 4).swapaxes(1,2).float()\n",
    "test_eq(PoolingLayer('conv1d', seq_len=4, token=False, seq_last=False)(t).shape, (2, 3))\n",
    "test_eq(PoolingLayer('max', token=False, seq_last=False)(t), tensor([[ 3.,  7., 11.], [15., 19., 23.]]))\n",
    "test_close(PoolingLayer('mean', token=False, seq_last=False)(t), tensor([[ 1.5000,  5.5000,  9.5000],\n",
    "                                                                        [13.5000, 17.5000, 21.5000]]))\n",
    "test_close(PoolingLayer('max-mean', token=False, seq_last=False)(t), tensor([[ 3.,  7., 11.,  1.5000,  5.5000,  9.5000],\n",
    "                                                                            [15., 19., 23., 13.5000, 17.5000, 21.5000]]))\n",
    "test_close(PoolingLayer('flatten', token=False, seq_last=False)(t), tensor([[ 0.,  4.,  8.,  1.,  5.,  9.,  2.,  6., 10.,  3.,  7., 11.],\n",
    "                                                                            [12., 16., 20., 13., 17., 21., 14., 18., 22., 15., 19., 23.]]))\n",
    "test_eq(PoolingLayer('conv1d', seq_len=4, token=False, seq_last=False)(t).shape, (2, 3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class GEGLU(Module):\n",
    "    def forward(self, x):\n",
    "        x, gates = x.chunk(2, dim=-1)\n",
    "        return x * F.gelu(gates)\n",
    "\n",
    "class ReGLU(Module):\n",
    "    def forward(self, x):\n",
    "        x, gates = x.chunk(2, dim=-1)\n",
    "        return x * F.relu(gates)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "pytorch_acts = [nn.ELU, nn.LeakyReLU, nn.PReLU, nn.ReLU, nn.ReLU6, nn.SELU, nn.CELU, nn.GELU, nn.Sigmoid, Mish, nn.Softplus,\n",
    "nn.Tanh, nn.Softmax, GEGLU, ReGLU, SmeLU]\n",
    "pytorch_act_names = [a.__name__.lower() for a in pytorch_acts]\n",
    "\n",
    "def get_act_fn(act, **act_kwargs):\n",
    "    if act is None: return\n",
    "    elif isinstance(act, nn.Module): return act\n",
    "    elif callable(act): return act(**act_kwargs)\n",
    "    idx = pytorch_act_names.index(act.lower())\n",
    "    return pytorch_acts[idx](**act_kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(get_act_fn(nn.ReLU).__repr__(), \"ReLU()\")\n",
    "test_eq(get_act_fn(nn.ReLU()).__repr__(), \"ReLU()\")\n",
    "test_eq(get_act_fn(nn.LeakyReLU, negative_slope=0.05).__repr__(), \"LeakyReLU(negative_slope=0.05)\")\n",
    "test_eq(get_act_fn('reglu').__repr__(), \"ReGLU()\")\n",
    "test_eq(get_act_fn('leakyrelu', negative_slope=0.05).__repr__(), \"LeakyReLU(negative_slope=0.05)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class RevIN(nn.Module):\n",
    "    \"\"\" Reversible Instance Normalization layer adapted from\n",
    "\n",
    "        Kim, T., Kim, J., Tae, Y., Park, C., Choi, J. H., & Choo, J. (2021, September).\n",
    "        Reversible instance normalization for accurate time-series forecasting against distribution shift.\n",
    "        In International Conference on Learning Representations.\n",
    "        Original code: https://github.com/ts-kim/RevIN\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self,\n",
    "         c_in:int,    # #features (aka variables or channels)\n",
    "         affine:bool=True,  # flag to incidate if RevIN has learnable weight and bias\n",
    "         subtract_last:bool=False,\n",
    "         dim:int=2,   # int or tuple of dimensions used to calculate mean and std\n",
    "         eps:float=1e-5  # epsilon - parameter added for numerical stability\n",
    "         ):\n",
    "        super().__init__()\n",
    "        self.c_in, self.affine, self.subtract_last, self.dim, self.eps = c_in, affine, subtract_last, dim, eps\n",
    "        if self.affine:\n",
    "            self.weight = nn.Parameter(torch.ones(1, c_in, 1))\n",
    "            self.bias = nn.Parameter(torch.zeros(1, c_in, 1))\n",
    "\n",
    "    def forward(self, x:Tensor, mode:Tensor):\n",
    "        \"\"\"Args:\n",
    "\n",
    "            x: rank 3 tensor with shape [batch size x c_in x sequence length]\n",
    "            mode: torch.tensor(True) to normalize data and torch.tensor(False) to reverse normalization\n",
    "        \"\"\"\n",
    "\n",
    "        # Normalize\n",
    "        if mode: return self.normalize(x)\n",
    "\n",
    "        # Denormalize\n",
    "        else: return self.denormalize(x)\n",
    "\n",
    "    def normalize(self, x):\n",
    "        if self.subtract_last:\n",
    "            self.sub = x[..., -1].unsqueeze(-1).detach()\n",
    "        else:\n",
    "            self.sub = torch.mean(x, dim=-1, keepdim=True).detach()\n",
    "        self.std = torch.std(x, dim=-1, keepdim=True, unbiased=False).detach() + self.eps\n",
    "        if self.affine:\n",
    "            x = x.sub(self.sub)\n",
    "            x = x.div(self.std)\n",
    "            x = x.mul(self.weight)\n",
    "            x = x.add(self.bias)\n",
    "            return x\n",
    "        else:\n",
    "            x = x.sub(self.sub)\n",
    "            x = x.div(self.std)\n",
    "            return x\n",
    "\n",
    "    def denormalize(self, x):\n",
    "        if self.affine:\n",
    "            x = x.sub(self.bias)\n",
    "            x = x.div(self.weight)\n",
    "            x = x.mul(self.std)\n",
    "            x = x.add(self.sub)\n",
    "            return x\n",
    "        else:\n",
    "            x = x.mul(self.std)\n",
    "            x = x.add(self.sub)\n",
    "            return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class RevIN(nn.Module):\n",
    "    \"\"\" Reversible Instance Normalization layer adapted from\n",
    "\n",
    "        Kim, T., Kim, J., Tae, Y., Park, C., Choi, J. H., & Choo, J. (2021, September).\n",
    "        Reversible instance normalization for accurate time-series forecasting against distribution shift.\n",
    "        In International Conference on Learning Representations.\n",
    "        Original code: https://github.com/ts-kim/RevIN\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self,\n",
    "         c_in:int,    # #features (aka variables or channels)\n",
    "         affine:bool=True,  # flag to incidate if RevIN has learnable weight and bias\n",
    "         subtract_last:bool=False,\n",
    "         dim:int=2,   # int or tuple of dimensions used to calculate mean and std\n",
    "         eps:float=1e-5  # epsilon - parameter added for numerical stability\n",
    "         ):\n",
    "        super().__init__()\n",
    "        self.c_in, self.affine, self.subtract_last, self.dim, self.eps = c_in, affine, subtract_last, dim, eps\n",
    "        self.weight = nn.Parameter(torch.ones(1, c_in, 1))\n",
    "        self.bias = nn.Parameter(torch.zeros(1, c_in, 1))\n",
    "        self.sub, self.std, self.mul, self.add = torch.zeros(1), torch.ones(1), torch.ones(1), torch.zeros(1)\n",
    "\n",
    "    def forward(self, x:Tensor, mode:Tensor):\n",
    "        \"\"\"Args:\n",
    "\n",
    "            x: rank 3 tensor with shape [batch size x c_in x sequence length]\n",
    "            mode: torch.tensor(True) to normalize data and torch.tensor(False) to reverse normalization\n",
    "        \"\"\"\n",
    "\n",
    "        # Normalize\n",
    "        if mode:\n",
    "            if self.subtract_last:\n",
    "                self.sub = x[..., -1].unsqueeze(-1).detach()\n",
    "            else:\n",
    "                self.sub = torch.mean(x, dim=-1, keepdim=True).detach()\n",
    "            self.std = torch.std(x, dim=-1, keepdim=True, unbiased=False).detach() + self.eps\n",
    "            if self.affine:\n",
    "                x = x.sub(self.sub)\n",
    "                x = x.div(self.std)\n",
    "                x = x.mul(self.weight)\n",
    "                x = x.add(self.bias)\n",
    "                return x\n",
    "            else:\n",
    "                x = x.sub(self.sub)\n",
    "                x = x.div(self.std)\n",
    "                return x\n",
    "\n",
    "        # Denormalize\n",
    "        else:\n",
    "            if self.affine:\n",
    "                x = x.sub(self.bias)\n",
    "                x = x.div(self.weight)\n",
    "                x = x.mul(self.std)\n",
    "                x = x.add(self.sub)\n",
    "                return x\n",
    "            else:\n",
    "                x = x.mul(self.std)\n",
    "                x = x.add(self.sub)\n",
    "                return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = ((torch.rand(16, 5, 100) - .25) * torch.Tensor([.01, .1, 1, 10, 100]).reshape(1, -1, 1)).cumsum(-1)\n",
    "t_clone = t.clone()\n",
    "l = RevIN(5)\n",
    "t_norm = l(t, torch.tensor(True))\n",
    "t_denorm = l(t_norm, torch.tensor(False))\n",
    "test_close(t_clone, t_denorm, eps=1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "scripting ok\n"
     ]
    }
   ],
   "source": [
    "model = RevIN(5, affine=True)\n",
    "try:\n",
    "    scripted_model = torch.jit.script(model)\n",
    "    file_path = f\"test_scripted_model.pt\"\n",
    "    torch.jit.save(scripted_model, file_path)\n",
    "    scripted_model = torch.jit.load(file_path)\n",
    "\n",
    "    inp = ((torch.rand(16, 5, 100) - .25) * torch.Tensor([.01, .1, 1, 10, 100]).reshape(1, -1, 1)).cumsum(-1)\n",
    "    normed_output = model(inp, torch.tensor(True))\n",
    "    demormed_output = model(normed_output, torch.tensor(False))\n",
    "    scripted_normed_output = scripted_model(inp, torch.tensor(True))\n",
    "    scripted_denormed_output = scripted_model(scripted_normed_output, torch.tensor(False))\n",
    "    test_close(normed_output, scripted_normed_output)\n",
    "    test_close(demormed_output, scripted_denormed_output)\n",
    "    os.remove(file_path)\n",
    "    del scripted_model\n",
    "    gc.collect()\n",
    "    print('scripting ok')\n",
    "except Exception as e:\n",
    "    print(f'scripting failed: {e}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def create_pool_head(n_in, c_out, seq_len=None, concat_pool=False, fc_dropout=0., bn=False, y_range=None, **kwargs):\n",
    "    if kwargs: print(f'{kwargs}  not being used')\n",
    "    if concat_pool: n_in*=2\n",
    "    layers = [GACP1d(1) if concat_pool else GAP1d(1)]\n",
    "    layers += [LinBnDrop(n_in, c_out, bn=bn, p=fc_dropout)]\n",
    "    if y_range: layers += [SigmoidRange(*y_range)]\n",
    "    return nn.Sequential(*layers)\n",
    "\n",
    "pool_head = create_pool_head\n",
    "average_pool_head = partial(pool_head, concat_pool=False)\n",
    "setattr(average_pool_head, \"__name__\", \"average_pool_head\")\n",
    "concat_pool_head = partial(pool_head, concat_pool=True)\n",
    "setattr(concat_pool_head, \"__name__\", \"concat_pool_head\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (0): GACP1d(\n",
       "    (gacp): AdaptiveConcatPool1d(\n",
       "      (ap): AdaptiveAvgPool1d(output_size=1)\n",
       "      (mp): AdaptiveMaxPool1d(output_size=1)\n",
       "    )\n",
       "    (flatten): Reshape(bs)\n",
       "  )\n",
       "  (1): LinBnDrop(\n",
       "    (0): BatchNorm1d(24, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (1): Dropout(p=0.5, inplace=False)\n",
       "    (2): Linear(in_features=24, out_features=2, bias=False)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 16\n",
    "nf = 12\n",
    "c_out = 2\n",
    "seq_len = 20\n",
    "t = torch.rand(bs, nf, seq_len)\n",
    "test_eq(create_pool_head(nf, c_out, seq_len, fc_dropout=0.5)(t).shape, (bs, c_out))\n",
    "test_eq(create_pool_head(nf, c_out, seq_len, concat_pool=True, fc_dropout=0.5)(t).shape, (bs, c_out))\n",
    "create_pool_head(nf, c_out, seq_len, concat_pool=True, bn=True, fc_dropout=.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def max_pool_head(n_in, c_out, seq_len, fc_dropout=0., bn=False, y_range=None, **kwargs):\n",
    "    if kwargs: print(f'{kwargs}  not being used')\n",
    "    layers = [nn.MaxPool1d(seq_len, **kwargs), Reshape()]\n",
    "    layers += [LinBnDrop(n_in, c_out, bn=bn, p=fc_dropout)]\n",
    "    if y_range: layers += [SigmoidRange(*y_range)]\n",
    "    return nn.Sequential(*layers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bs = 16\n",
    "nf = 12\n",
    "c_out = 2\n",
    "seq_len = 20\n",
    "t = torch.rand(bs, nf, seq_len)\n",
    "test_eq(max_pool_head(nf, c_out, seq_len, fc_dropout=0.5)(t).shape, (bs, c_out))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def create_pool_plus_head(*args, lin_ftrs=None, fc_dropout=0., concat_pool=True, bn_final=False, lin_first=False, y_range=None):\n",
    "    nf = args[0]\n",
    "    c_out = args[1]\n",
    "    if concat_pool: nf = nf * 2\n",
    "    lin_ftrs = [nf, 512, c_out] if lin_ftrs is None else [nf] + lin_ftrs + [c_out]\n",
    "    ps = L(fc_dropout)\n",
    "    if len(ps) == 1: ps = [ps[0]/2] * (len(lin_ftrs)-2) + ps\n",
    "    actns = [nn.ReLU(inplace=True)] * (len(lin_ftrs)-2) + [None]\n",
    "    pool = AdaptiveConcatPool1d() if concat_pool else nn.AdaptiveAvgPool1d(1)\n",
    "    layers = [pool, Reshape()]\n",
    "    if lin_first: layers.append(nn.Dropout(ps.pop(0)))\n",
    "    for ni,no,p,actn in zip(lin_ftrs[:-1], lin_ftrs[1:], ps, actns):\n",
    "        layers += LinBnDrop(ni, no, bn=True, p=p, act=actn, lin_first=lin_first)\n",
    "    if lin_first: layers.append(nn.Linear(lin_ftrs[-2], c_out))\n",
    "    if bn_final: layers.append(nn.BatchNorm1d(lin_ftrs[-1], momentum=0.01))\n",
    "    if y_range is not None: layers.append(SigmoidRange(*y_range))\n",
    "    return nn.Sequential(*layers)\n",
    "\n",
    "pool_plus_head = create_pool_plus_head"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (0): AdaptiveConcatPool1d(\n",
       "    (ap): AdaptiveAvgPool1d(output_size=1)\n",
       "    (mp): AdaptiveMaxPool1d(output_size=1)\n",
       "  )\n",
       "  (1): Reshape(bs)\n",
       "  (2): BatchNorm1d(24, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (3): Dropout(p=0.25, inplace=False)\n",
       "  (4): Linear(in_features=24, out_features=512, bias=False)\n",
       "  (5): ReLU(inplace=True)\n",
       "  (6): BatchNorm1d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (7): Dropout(p=0.5, inplace=False)\n",
       "  (8): Linear(in_features=512, out_features=2, bias=False)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 16\n",
    "nf = 12\n",
    "c_out = 2\n",
    "seq_len = 20\n",
    "t = torch.rand(bs, nf, seq_len)\n",
    "test_eq(create_pool_plus_head(nf, c_out, seq_len, fc_dropout=0.5)(t).shape, (bs, c_out))\n",
    "test_eq(create_pool_plus_head(nf, c_out, concat_pool=True, fc_dropout=0.5)(t).shape, (bs, c_out))\n",
    "create_pool_plus_head(nf, c_out, seq_len, fc_dropout=0.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def create_conv_head(*args, adaptive_size=None, y_range=None):\n",
    "    nf = args[0]\n",
    "    c_out = args[1]\n",
    "    layers = [nn.AdaptiveAvgPool1d(adaptive_size)] if adaptive_size is not None else []\n",
    "    for i in range(2):\n",
    "        if nf > 1:\n",
    "            layers += [ConvBlock(nf, nf // 2, 1)]\n",
    "            nf = nf//2\n",
    "        else: break\n",
    "    layers += [ConvBlock(nf, c_out, 1), GAP1d(1)]\n",
    "    if y_range: layers += [SigmoidRange(*y_range)]\n",
    "    return nn.Sequential(*layers)\n",
    "\n",
    "conv_head = create_conv_head"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (0): ConvBlock(\n",
       "    (0): Conv1d(12, 6, kernel_size=(1,), stride=(1,), bias=False)\n",
       "    (1): BatchNorm1d(6, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (2): ReLU()\n",
       "  )\n",
       "  (1): ConvBlock(\n",
       "    (0): Conv1d(6, 3, kernel_size=(1,), stride=(1,), bias=False)\n",
       "    (1): BatchNorm1d(3, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (2): ReLU()\n",
       "  )\n",
       "  (2): ConvBlock(\n",
       "    (0): Conv1d(3, 2, kernel_size=(1,), stride=(1,), bias=False)\n",
       "    (1): BatchNorm1d(2, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (2): ReLU()\n",
       "  )\n",
       "  (3): GAP1d(\n",
       "    (gap): AdaptiveAvgPool1d(output_size=1)\n",
       "    (flatten): Reshape(bs)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 16\n",
    "nf = 12\n",
    "c_out = 2\n",
    "seq_len = 20\n",
    "t = torch.rand(bs, nf, seq_len)\n",
    "test_eq(create_conv_head(nf, c_out, seq_len)(t).shape, (bs, c_out))\n",
    "test_eq(create_conv_head(nf, c_out, adaptive_size=50)(t).shape, (bs, c_out))\n",
    "create_conv_head(nf, c_out, 50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def create_mlp_head(nf, c_out, seq_len=None, flatten=True, fc_dropout=0., bn=False, lin_first=False, y_range=None):\n",
    "    if flatten: nf *= seq_len\n",
    "    layers = [Reshape()] if flatten else []\n",
    "    layers += [LinBnDrop(nf, c_out, bn=bn, p=fc_dropout, lin_first=lin_first)]\n",
    "    if y_range: layers += [SigmoidRange(*y_range)]\n",
    "    return nn.Sequential(*layers)\n",
    "\n",
    "mlp_head = create_mlp_head"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (0): Reshape(bs)\n",
       "  (1): LinBnDrop(\n",
       "    (0): BatchNorm1d(240, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (1): Dropout(p=0.5, inplace=False)\n",
       "    (2): Linear(in_features=240, out_features=2, bias=False)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 16\n",
    "nf = 12\n",
    "c_out = 2\n",
    "seq_len = 20\n",
    "t = torch.rand(bs, nf, seq_len)\n",
    "test_eq(create_mlp_head(nf, c_out, seq_len, fc_dropout=0.5)(t).shape, (bs, c_out))\n",
    "t = torch.rand(bs, nf, seq_len)\n",
    "create_mlp_head(nf, c_out, seq_len, bn=True, fc_dropout=.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def create_fc_head(nf, c_out, seq_len=None, flatten=True, lin_ftrs=None, y_range=None, fc_dropout=0., bn=False, bn_final=False, act=nn.ReLU(inplace=True)):\n",
    "    if flatten: nf *= seq_len\n",
    "    layers = [Reshape()] if flatten else []\n",
    "    lin_ftrs = [nf, 512, c_out] if lin_ftrs is None else [nf] + lin_ftrs + [c_out]\n",
    "    if not is_listy(fc_dropout): fc_dropout = [fc_dropout]*(len(lin_ftrs) - 1)\n",
    "    actns = [act for _ in range(len(lin_ftrs) - 2)] + [None]\n",
    "    layers += [LinBnDrop(lin_ftrs[i], lin_ftrs[i+1], bn=bn and (i!=len(actns)-1 or bn_final), p=p, act=a) for i,(p,a) in enumerate(zip(fc_dropout+[0.], actns))]\n",
    "    if y_range is not None: layers.append(SigmoidRange(*y_range))\n",
    "    return nn.Sequential(*layers)\n",
    "\n",
    "fc_head = create_fc_head"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (0): Reshape(bs)\n",
       "  (1): LinBnDrop(\n",
       "    (0): BatchNorm1d(240, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (1): Dropout(p=0.5, inplace=False)\n",
       "    (2): Linear(in_features=240, out_features=2, bias=False)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 16\n",
    "nf = 12\n",
    "c_out = 2\n",
    "seq_len = 20\n",
    "t = torch.rand(bs, nf, seq_len)\n",
    "test_eq(create_fc_head(nf, c_out, seq_len, fc_dropout=0.5)(t).shape, (bs, c_out))\n",
    "create_mlp_head(nf, c_out, seq_len, bn=True, fc_dropout=.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def create_rnn_head(*args, fc_dropout=0., bn=False, y_range=None):\n",
    "    nf = args[0]\n",
    "    c_out = args[1]\n",
    "    layers = [LastStep()]\n",
    "    layers += [LinBnDrop(nf, c_out, bn=bn, p=fc_dropout)]\n",
    "    if y_range: layers += [SigmoidRange(*y_range)]\n",
    "    return nn.Sequential(*layers)\n",
    "\n",
    "rnn_head = create_rnn_head"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (0): LastStep()\n",
       "  (1): LinBnDrop(\n",
       "    (0): BatchNorm1d(12, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (1): Dropout(p=0.5, inplace=False)\n",
       "    (2): Linear(in_features=12, out_features=2, bias=False)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 16\n",
    "nf = 12\n",
    "c_out = 2\n",
    "seq_len = 20\n",
    "t = torch.rand(bs, nf, seq_len)\n",
    "test_eq(create_rnn_head(nf, c_out, seq_len, fc_dropout=0.5)(t).shape, (bs, c_out))\n",
    "create_rnn_head(nf, c_out, seq_len, bn=True, fc_dropout=.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def imputation_head(c_in, c_out, seq_len=None, ks=1, y_range=None, fc_dropout=0.):\n",
    "    layers = [nn.Dropout(fc_dropout), nn.Conv1d(c_in, c_out, ks)]\n",
    "    if y_range is not None:\n",
    "        y_range = (tensor(y_range[0]), tensor(y_range[1]))\n",
    "        layers += [SigmoidRange(*y_range)]\n",
    "    return nn.Sequential(*layers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (0): Dropout(p=0.0, inplace=False)\n",
       "  (1): Conv1d(12, 2, kernel_size=(1,), stride=(1,))\n",
       "  (2): fastai.layers.SigmoidRange(low=tensor([0.1000, 0.1000, 0.1000, 0.1000, 0.2000, 0.2000, 0.2000, 0.2000, 0.3000,\n",
       "          0.3000, 0.3000, 0.3000]), high=tensor([0.6000, 0.6000, 0.6000, 0.6000, 0.7000, 0.7000, 0.7000, 0.7000, 0.8000,\n",
       "          0.8000, 0.8000, 0.8000]))\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 16\n",
    "nf = 12\n",
    "ni = 2\n",
    "seq_len = 20\n",
    "t = torch.rand(bs, nf, seq_len)\n",
    "head = imputation_head(nf, ni, seq_len=None, ks=1, y_range=None, fc_dropout=0.)\n",
    "test_eq(head(t).shape, (bs, ni, seq_len))\n",
    "head = imputation_head(nf, ni, seq_len=None, ks=1, y_range=(.3,.7), fc_dropout=0.)\n",
    "test_ge(head(t).min(), .3)\n",
    "test_le(head(t).max(), .7)\n",
    "y_range = (tensor([0.1000, 0.1000, 0.1000, 0.1000, 0.2000, 0.2000, 0.2000, 0.2000, 0.3000,\n",
    "                   0.3000, 0.3000, 0.3000]),\n",
    "           tensor([0.6000, 0.6000, 0.6000, 0.6000, 0.7000, 0.7000, 0.7000, 0.7000, 0.8000,\n",
    "                   0.8000, 0.8000, 0.8000]))\n",
    "test_ge(head(t).min(), .1)\n",
    "test_le(head(t).max(), .9)\n",
    "head = imputation_head(nf, ni, seq_len=None, ks=1, y_range=y_range, fc_dropout=0.)\n",
    "head"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class create_conv_lin_nd_head(nn.Sequential):\n",
    "    \"Module to create a nd output head\"\n",
    "\n",
    "    def __init__(self, n_in, n_out, seq_len, d, conv_first=True, conv_bn=False, lin_bn=False, fc_dropout=0., **kwargs):\n",
    "\n",
    "        assert d, \"you cannot use an nd head when d is None or 0\"\n",
    "        if is_listy(d):\n",
    "            fd = 1\n",
    "            shape = []\n",
    "            for _d in d:\n",
    "                fd *= _d\n",
    "                shape.append(_d)\n",
    "            if n_out > 1: shape.append(n_out)\n",
    "        else:\n",
    "            fd = d\n",
    "            shape = [d, n_out] if n_out > 1 else [d]\n",
    "\n",
    "        conv = [BatchNorm(n_in, ndim=1)] if conv_bn else []\n",
    "        conv.append(Conv1d(n_in, n_out, 1, padding=0, bias=not conv_bn, **kwargs))\n",
    "        l = [Transpose(-1, -2), BatchNorm(seq_len, ndim=1), Transpose(-1, -2)] if lin_bn else []\n",
    "        if fc_dropout != 0: l.append(nn.Dropout(fc_dropout))\n",
    "        lin = [nn.Linear(seq_len, fd, bias=not lin_bn)]\n",
    "        lin_layers = l+lin\n",
    "        layers = conv + lin_layers if conv_first else lin_layers + conv\n",
    "        layers += [Transpose(-1,-2)]\n",
    "        layers += [Reshape(*shape)]\n",
    "\n",
    "        super().__init__(*layers)\n",
    "\n",
    "conv_lin_nd_head = create_conv_lin_nd_head\n",
    "conv_lin_3d_head = create_conv_lin_nd_head # included for compatibility\n",
    "create_conv_lin_3d_head = create_conv_lin_nd_head # included for compatibility"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(TensorBase(1.7252, grad_fn=<AliasBackward0>),\n",
       " create_conv_lin_nd_head(\n",
       "   (0): Conv1d(32, 5, kernel_size=(1,), stride=(1,))\n",
       "   (1): Dropout(p=0.5, inplace=False)\n",
       "   (2): Linear(in_features=10, out_features=2, bias=True)\n",
       "   (3): Transpose(dims=-1, -2).contiguous()\n",
       "   (4): Reshape(bs, 2, 5)\n",
       " ))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 16\n",
    "nf = 32\n",
    "c = 5\n",
    "seq_len = 10\n",
    "d = 2\n",
    "targ = torch.randint(0, c, (bs,d))\n",
    "t = torch.randn(bs, nf, seq_len)\n",
    "head = conv_lin_nd_head(nf, c, seq_len, d, conv_first=True, fc_dropout=.5)\n",
    "inp = head(t)\n",
    "test_eq(inp.shape, (bs, d, c))\n",
    "loss = CrossEntropyLossFlat()(inp, targ)\n",
    "loss, head"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(TensorBase(1.6647, grad_fn=<AliasBackward0>),\n",
       " create_conv_lin_nd_head(\n",
       "   (0): Dropout(p=0.5, inplace=False)\n",
       "   (1): Linear(in_features=10, out_features=16, bias=True)\n",
       "   (2): Conv1d(32, 5, kernel_size=(1,), stride=(1,))\n",
       "   (3): Transpose(dims=-1, -2).contiguous()\n",
       "   (4): Reshape(bs, 2, 8, 5)\n",
       " ))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 16\n",
    "nf = 32\n",
    "c = 5\n",
    "seq_len = 10\n",
    "d = [2, 8]\n",
    "targ = torch.randint(0, c, [bs]+d)\n",
    "t = torch.randn(bs, nf, seq_len)\n",
    "head = conv_lin_nd_head(nf, c, seq_len, d, conv_first=False, fc_dropout=.5)\n",
    "inp = head(t)\n",
    "test_eq(inp.shape, [bs]+d+[c])\n",
    "loss = CrossEntropyLossFlat()(inp, targ)\n",
    "loss, head"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(TensorBase(0.7063, grad_fn=<AliasBackward0>),\n",
       " create_conv_lin_nd_head(\n",
       "   (0): Dropout(p=0.5, inplace=False)\n",
       "   (1): Linear(in_features=10, out_features=2, bias=True)\n",
       "   (2): Conv1d(32, 1, kernel_size=(1,), stride=(1,))\n",
       "   (3): Transpose(dims=-1, -2).contiguous()\n",
       "   (4): Reshape(bs, 2)\n",
       " ))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 16\n",
    "nf = 32\n",
    "c = 1\n",
    "seq_len = 10\n",
    "d = 2\n",
    "targ = torch.rand(bs, d)\n",
    "t = torch.randn(bs, nf, seq_len)\n",
    "head = conv_lin_nd_head(nf, c, seq_len, d, conv_first=False, fc_dropout=.5)\n",
    "inp = head(t)\n",
    "test_eq(inp.shape, (bs, d))\n",
    "loss = L1LossFlat()(inp, targ)\n",
    "loss, head"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(TensorBase(0.6203, grad_fn=<AliasBackward0>),\n",
       " create_conv_lin_nd_head(\n",
       "   (0): Dropout(p=0.5, inplace=False)\n",
       "   (1): Linear(in_features=10, out_features=6, bias=True)\n",
       "   (2): Conv1d(32, 1, kernel_size=(1,), stride=(1,))\n",
       "   (3): Transpose(dims=-1, -2).contiguous()\n",
       "   (4): Reshape(bs, 2, 3)\n",
       " ))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 16\n",
    "nf = 32\n",
    "c = 1\n",
    "seq_len = 10\n",
    "d = [2,3]\n",
    "targ = torch.rand(bs, *d)\n",
    "t = torch.randn(bs, nf, seq_len)\n",
    "head = conv_lin_nd_head(nf, c, seq_len, d, conv_first=False, fc_dropout=.5)\n",
    "inp = head(t)\n",
    "test_eq(inp.shape, [bs]+d)\n",
    "loss = L1LossFlat()(inp, targ)\n",
    "loss, head"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class lin_nd_head(nn.Sequential):\n",
    "    \"Module to create a nd output head with linear layers\"\n",
    "\n",
    "    def __init__(self, n_in, n_out, seq_len=None, d=None, flatten=False, use_bn=False, fc_dropout=0.):\n",
    "\n",
    "        if seq_len is None:\n",
    "            seq_len = 1\n",
    "        if d is None:\n",
    "            fd = 1\n",
    "            shape = [n_out]\n",
    "        elif is_listy(d):\n",
    "            fd = 1\n",
    "            shape = []\n",
    "            for _d in d:\n",
    "                fd *= _d\n",
    "                shape.append(_d)\n",
    "            if n_out > 1: shape.append(n_out)\n",
    "        else:\n",
    "            fd = d\n",
    "            shape = [d, n_out] if n_out > 1 else [d]\n",
    "\n",
    "        layers = []\n",
    "        if use_bn:\n",
    "            layers += [nn.BatchNorm1d(n_in)]\n",
    "        if fc_dropout:\n",
    "            layers += [nn.Dropout(fc_dropout)]\n",
    "        if d is None:\n",
    "            if not flatten or seq_len == 1:\n",
    "                layers += [nn.AdaptiveAvgPool1d(1), Squeeze(-1), nn.Linear(n_in, n_out)]\n",
    "                if n_out == 1:\n",
    "                    layers += [Squeeze(-1)]\n",
    "            else:\n",
    "                layers += [Reshape(), nn.Linear(n_in * seq_len, n_out * fd)]\n",
    "                if n_out * fd== 1:\n",
    "                    layers += [Squeeze(-1)]\n",
    "        else:\n",
    "            if seq_len == 1:\n",
    "                layers += [nn.AdaptiveAvgPool1d(1)]\n",
    "            if not flatten and fd == seq_len:\n",
    "                layers += [Transpose(1,2), nn.Linear(n_in, n_out)]\n",
    "            else:\n",
    "                layers += [Reshape(), nn.Linear(n_in * seq_len, n_out * fd)]\n",
    "            layers += [Reshape(*shape)]\n",
    "\n",
    "        super().__init__(*layers)\n",
    "\n",
    "create_lin_nd_head = lin_nd_head\n",
    "lin_3d_head = lin_nd_head # included for backwards compatiblity\n",
    "create_lin_3d_head = lin_nd_head # included for backwards compatiblity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bs = 16\n",
    "nf = 32\n",
    "seq_len = 50\n",
    "x = torch.normal(0, 1, (bs, nf, seq_len))\n",
    "\n",
    "for use_bn in [False, True]:\n",
    "    for fc_dropout in [0, 0.2]:\n",
    "        for flatten in [False, True]:\n",
    "            for c in [1, 3]:\n",
    "                for d in [None, (50,), (50,10), (30,5), (50,2,3), (30,2,3)]:\n",
    "                    for q_len in [1, seq_len]:\n",
    "                        head = lin_nd_head(nf, c, q_len, d, flatten=flatten, use_bn=use_bn, fc_dropout=fc_dropout)\n",
    "                        test_eq(head(x).shape, (bs, ) + (d if d is not None else ()) + ((c,) if c != 1 else ()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(TensorBase(1.7711, grad_fn=<AliasBackward0>),\n",
       " lin_nd_head(\n",
       "   (0): Dropout(p=0.5, inplace=False)\n",
       "   (1): Reshape(bs)\n",
       "   (2): Linear(in_features=320, out_features=10, bias=True)\n",
       "   (3): Reshape(bs, 2, 5)\n",
       " ))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 16\n",
    "nf = 32\n",
    "c = 5\n",
    "seq_len = 10\n",
    "d = 2\n",
    "targ = torch.randint(0, c, (bs,d))\n",
    "t = torch.randn(bs, nf, seq_len)\n",
    "head = lin_nd_head(nf, c, seq_len, d, fc_dropout=.5)\n",
    "inp = head(t)\n",
    "test_eq(inp.shape, (bs, d, c))\n",
    "loss = CrossEntropyLossFlat()(inp, targ)\n",
    "loss, head"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(TensorBase(1.8884, grad_fn=<AliasBackward0>),\n",
       " lin_nd_head(\n",
       "   (0): Dropout(p=0.5, inplace=False)\n",
       "   (1): Reshape(bs)\n",
       "   (2): Linear(in_features=320, out_features=80, bias=True)\n",
       "   (3): Reshape(bs, 2, 8, 5)\n",
       " ))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 16\n",
    "nf = 32\n",
    "c = 5\n",
    "seq_len = 10\n",
    "d = [2, 8]\n",
    "targ = torch.randint(0, c, [bs]+d)\n",
    "t = torch.randn(bs, nf, seq_len)\n",
    "head = lin_nd_head(nf, c, seq_len, d, fc_dropout=.5)\n",
    "inp = head(t)\n",
    "test_eq(inp.shape, [bs]+d+[c])\n",
    "loss = CrossEntropyLossFlat()(inp, targ)\n",
    "loss, head"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(TensorBase(0.7737, grad_fn=<AliasBackward0>),\n",
       " lin_nd_head(\n",
       "   (0): Dropout(p=0.5, inplace=False)\n",
       "   (1): Reshape(bs)\n",
       "   (2): Linear(in_features=320, out_features=2, bias=True)\n",
       "   (3): Reshape(bs, 2)\n",
       " ))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 16\n",
    "nf = 32\n",
    "c = 1\n",
    "seq_len = 10\n",
    "d = 2\n",
    "targ = torch.rand(bs, d)\n",
    "t = torch.randn(bs, nf, seq_len)\n",
    "head = lin_nd_head(nf, c, seq_len, d, fc_dropout=.5)\n",
    "inp = head(t)\n",
    "test_eq(inp.shape, (bs, d))\n",
    "loss = L1LossFlat()(inp, targ)\n",
    "loss, head"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(TensorBase(0.8873, grad_fn=<AliasBackward0>),\n",
       " lin_nd_head(\n",
       "   (0): Dropout(p=0.5, inplace=False)\n",
       "   (1): Reshape(bs)\n",
       "   (2): Linear(in_features=320, out_features=6, bias=True)\n",
       "   (3): Reshape(bs, 2, 3)\n",
       " ))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 16\n",
    "nf = 32\n",
    "c = 1\n",
    "seq_len = 10\n",
    "d = [2,3]\n",
    "targ = torch.rand(bs, *d)\n",
    "t = torch.randn(bs, nf, seq_len)\n",
    "head = lin_nd_head(nf, c, seq_len, d, fc_dropout=.5)\n",
    "inp = head(t)\n",
    "test_eq(inp.shape, [bs]+d)\n",
    "loss = L1LossFlat()(inp, targ)\n",
    "loss, head"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class rocket_nd_head(nn.Sequential):\n",
    "    \"Module to create a nd output head with linear layers for the rocket family of models\"\n",
    "\n",
    "    def __init__(self, n_in, n_out, seq_len=None, d=None, use_bn=False, fc_dropout=0., zero_init=True):\n",
    "\n",
    "        if d is None:\n",
    "            fd = 1\n",
    "            shape = [n_out]\n",
    "        elif is_listy(d):\n",
    "            fd = 1\n",
    "            shape = []\n",
    "            for _d in d:\n",
    "                fd *= _d\n",
    "                shape.append(_d)\n",
    "            if n_out > 1: shape.append(n_out)\n",
    "        else:\n",
    "            fd = d\n",
    "            shape = [d, n_out] if n_out > 1 else [d]\n",
    "\n",
    "        layers = [nn.Flatten()]\n",
    "        if use_bn:\n",
    "            layers += [nn.BatchNorm1d(n_in)]\n",
    "        if fc_dropout:\n",
    "            layers += [nn.Dropout(fc_dropout)]\n",
    "        linear = nn.Linear(n_in, fd * n_out)\n",
    "        if zero_init:\n",
    "            nn.init.constant_(linear.weight.data, 0)\n",
    "            nn.init.constant_(linear.bias.data, 0)\n",
    "        layers += [linear]\n",
    "        if d is None and n_out == 1:\n",
    "            layers += [Squeeze(-1)]\n",
    "        if d is not None:\n",
    "            layers += [Reshape(*shape)]\n",
    "\n",
    "        super().__init__(*layers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bs = 16\n",
    "nf = 99\n",
    "seq_len = 1\n",
    "x = torch.normal(0, 1, (bs, nf, seq_len))\n",
    "\n",
    "for use_bn in [False, True]:\n",
    "    for fc_dropout in [0, 0.2]:\n",
    "        for c in [1, 3]:\n",
    "            for d in [None, (50,), (50,10), (30,5), (50,2,3), (30,2,3)]:\n",
    "                head = rocket_nd_head(nf, c, 1, d, use_bn=use_bn, fc_dropout=fc_dropout)\n",
    "                test_eq(head(x).shape, (bs, ) + (d if d is not None else ()) + ((c,) if c != 1 else ()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class xresnet1d_nd_head(nn.Sequential):\n",
    "    \"Module to create a nd output head with linear layers for the xresnet family of models\"\n",
    "\n",
    "    def __init__(self, n_in, n_out, seq_len=None, d=None, use_bn=False, fc_dropout=0., zero_init=True):\n",
    "\n",
    "        if d is None:\n",
    "            fd = 1\n",
    "            shape = [n_out]\n",
    "        elif is_listy(d):\n",
    "            fd = 1\n",
    "            shape = []\n",
    "            for _d in d:\n",
    "                fd *= _d\n",
    "                shape.append(_d)\n",
    "            if n_out > 1: shape.append(n_out)\n",
    "        else:\n",
    "            fd = d\n",
    "            shape = [d, n_out] if n_out > 1 else [d]\n",
    "\n",
    "        layers = [nn.AdaptiveAvgPool1d(1), nn.Flatten()]\n",
    "        if use_bn:\n",
    "            layers += [nn.BatchNorm1d(n_in)]\n",
    "        if fc_dropout:\n",
    "            layers += [nn.Dropout(fc_dropout)]\n",
    "        linear = nn.Linear(n_in, fd * n_out)\n",
    "        if zero_init:\n",
    "            nn.init.constant_(linear.weight.data, 0)\n",
    "            nn.init.constant_(linear.bias.data, 0)\n",
    "        layers += [linear]\n",
    "        if d is None and n_out == 1:\n",
    "            layers += [Squeeze(-1)]\n",
    "        if d is not None:\n",
    "            layers += [Reshape(*shape)]\n",
    "\n",
    "        super().__init__(*layers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bs = 16\n",
    "nf = 99\n",
    "seq_len = 2\n",
    "x = torch.normal(0, 1, (bs, nf, seq_len))\n",
    "\n",
    "for use_bn in [False, True]:\n",
    "    for fc_dropout in [0, 0.2]:\n",
    "        for c in [1, 3]:\n",
    "            for d in [None, (50,), (50,10), (30,5), (50,2,3), (30,2,3)]:\n",
    "                head = xresnet1d_nd_head(nf, c, 1, d, use_bn=use_bn, fc_dropout=fc_dropout)\n",
    "                test_eq(head(x).shape, (bs, ) + (d if d is not None else ()) + ((c,) if c != 1 else ()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class create_conv_3d_head(nn.Sequential):\n",
    "    \"Module to create a nd output head with a convolutional layer\"\n",
    "    def __init__(self, n_in, n_out, seq_len, d, use_bn=False, **kwargs):\n",
    "        assert d, \"you cannot use an 3d head when d is None or 0\"\n",
    "        assert d == seq_len, 'You can only use this head when learn.dls.len == learn.dls.d'\n",
    "        layers = [nn.BatchNorm1d(n_in)] if use_bn else []\n",
    "        layers += [Conv(n_in, n_out, 1, **kwargs), Transpose(-1,-2)]\n",
    "        if n_out == 1: layers += [Squeeze(-1)]\n",
    "        super().__init__(*layers)\n",
    "\n",
    "conv_3d_head = create_conv_3d_head"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(TensorBase(1.8352, grad_fn=<AliasBackward0>),\n",
       " create_conv_3d_head(\n",
       "   (0): ConvBlock(\n",
       "     (0): Conv1d(32, 5, kernel_size=(1,), stride=(1,))\n",
       "   )\n",
       "   (1): Transpose(dims=-1, -2).contiguous()\n",
       " ))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 16\n",
    "nf = 32\n",
    "c = 5\n",
    "seq_len = 10\n",
    "d = 10\n",
    "targ = torch.randint(0, c, (bs,d))\n",
    "t = torch.randn(bs, nf, seq_len)\n",
    "head = conv_3d_head(nf, c, seq_len, d)\n",
    "inp = head(t)\n",
    "test_eq(inp.shape, (bs, d, c))\n",
    "loss = CrossEntropyLossFlat()(inp, targ)\n",
    "loss, head"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(TensorBase(0.6711, grad_fn=<AliasBackward0>),\n",
       " create_conv_3d_head(\n",
       "   (0): ConvBlock(\n",
       "     (0): Conv1d(32, 1, kernel_size=(1,), stride=(1,))\n",
       "   )\n",
       "   (1): Transpose(dims=-1, -2).contiguous()\n",
       "   (2): Squeeze(dim=-1)\n",
       " ))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 16\n",
    "nf = 32\n",
    "c = 1\n",
    "seq_len = 10\n",
    "d = 10\n",
    "targ = torch.rand(bs, d)\n",
    "t = torch.randn(bs, nf, seq_len)\n",
    "head = conv_3d_head(nf, c, seq_len, d)\n",
    "inp = head(t)\n",
    "test_eq(inp.shape, (bs, d))\n",
    "loss = L1LossFlat()(inp, targ)\n",
    "loss, head"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def universal_pool_head(n_in, c_out, seq_len, mult=2, pool_n_layers=2, pool_ln=True, pool_dropout=0.5, pool_act=nn.ReLU(),\n",
    "                        zero_init=True, bn=True, fc_dropout=0.):\n",
    "    return nn.Sequential(AdaptiveWeightedAvgPool1d(n_in, seq_len, n_layers=pool_n_layers, mult=mult, ln=pool_ln, dropout=pool_dropout, act=pool_act),\n",
    "                         Reshape(), LinBnDrop(n_in, c_out, p=fc_dropout, bn=bn))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bs, c_in, seq_len = 16, 128, 50\n",
    "c_out = 14\n",
    "t = torch.rand(bs, c_in, seq_len)\n",
    "uph = universal_pool_head(c_in, c_out, seq_len)\n",
    "test_eq(uph(t).shape, (bs, c_out))\n",
    "uph = universal_pool_head(c_in, c_out, seq_len, 2)\n",
    "test_eq(uph(t).shape, (bs, c_out))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "heads = [mlp_head, fc_head, average_pool_head, max_pool_head, concat_pool_head, pool_plus_head, conv_head, rnn_head,\n",
    "         conv_lin_nd_head, lin_nd_head, conv_3d_head, attentional_pool_head, universal_pool_head, gwa_pool_head]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "create_mlp_head\n",
      "create_fc_head\n",
      "average_pool_head\n",
      "max_pool_head\n",
      "concat_pool_head\n",
      "create_pool_plus_head\n",
      "create_conv_head\n",
      "create_rnn_head\n",
      "create_conv_lin_nd_head\n",
      "lin_nd_head\n",
      "create_conv_3d_head\n",
      "attentional_pool_head\n",
      "universal_pool_head\n",
      "gwa_pool_head\n"
     ]
    }
   ],
   "source": [
    "bs, c_in, seq_len = 16, 128, 50\n",
    "c_out = 14\n",
    "d = 5\n",
    "t = torch.rand(bs, c_in, seq_len)\n",
    "for head in heads:\n",
    "    print(head.__name__)\n",
    "    if head.__name__ == \"create_conv_3d_head\":\n",
    "        h = head(c_in, c_out, seq_len, seq_len)\n",
    "        test_eq(h(t).shape, (bs, seq_len, c_out))\n",
    "    elif 'nd' in head.__name__:\n",
    "        h = head(c_in, c_out, seq_len, d)\n",
    "        test_eq(h(t).shape, (bs, d, c_out))\n",
    "    else:\n",
    "        h = head(c_in, c_out, seq_len)\n",
    "        test_eq(h(t).shape, (bs, c_out))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class SqueezeExciteBlock(Module):\n",
    "    def __init__(self, ni, reduction=16):\n",
    "        self.avg_pool = GAP1d(1)\n",
    "        self.fc = nn.Sequential(nn.Linear(ni, ni // reduction, bias=False), nn.ReLU(),  nn.Linear(ni // reduction, ni, bias=False), nn.Sigmoid())\n",
    "\n",
    "    def forward(self, x):\n",
    "        y = self.avg_pool(x)\n",
    "        y = self.fc(y).unsqueeze(2)\n",
    "        return x * y.expand_as(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bs = 2\n",
    "ni = 32\n",
    "sl = 4\n",
    "t = torch.rand(bs, ni, sl)\n",
    "test_eq(SqueezeExciteBlock(ni)(t).shape, (bs, ni, sl))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class GaussianNoise(Module):\n",
    "    \"\"\"Gaussian noise regularizer.\n",
    "\n",
    "    Args:\n",
    "        sigma (float, optional): relative standard deviation used to generate the\n",
    "            noise. Relative means that it will be multiplied by the magnitude of\n",
    "            the value your are adding the noise to. This means that sigma can be\n",
    "            the same regardless of the scale of the vector.\n",
    "        is_relative_detach (bool, optional): whether to detach the variable before\n",
    "            computing the scale of the noise. If `False` then the scale of the noise\n",
    "            won't be seen as a constant but something to optimize: this will bias the\n",
    "            network to generate vectors with smaller values.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, sigma=0.1, is_relative_detach=True):\n",
    "        self.sigma, self.is_relative_detach = sigma, is_relative_detach\n",
    "\n",
    "    def forward(self, x):\n",
    "        if self.training and self.sigma not in [0, None]:\n",
    "            scale = self.sigma * (x.detach() if self.is_relative_detach else x)\n",
    "            sampled_noise = torch.empty(x.size(), device=x.device).normal_() * scale\n",
    "            x = x + sampled_noise\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = torch.ones(2,3,4)\n",
    "test_ne(GaussianNoise()(t), t)\n",
    "test_eq(GaussianNoise()(t).shape, t.shape)\n",
    "t = torch.ones(2,3)\n",
    "test_ne(GaussianNoise()(t), t)\n",
    "test_eq(GaussianNoise()(t).shape, t.shape)\n",
    "t = torch.ones(2)\n",
    "test_ne(GaussianNoise()(t), t)\n",
    "test_eq(GaussianNoise()(t).shape, t.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|eval: false\n",
    "#|hide\n",
    "import numpy as np\n",
    "from scipy.stats import ttest_ind, ttest_ind_from_stats\n",
    "from scipy.special import stdtr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ttest_ind:            t = -1.5827  p = 0.118873\n",
      "ttest_ind_from_stats: t = -1.5827  p = 0.118873\n",
      "formula:              t = -1.5827  p = 0.118873\n",
      "formula:              t = -1.5827\n"
     ]
    }
   ],
   "source": [
    "#|eval: false\n",
    "#|hide\n",
    "# https://stackoverflow.com/questions/22611446/perform-2-sample-t-test\n",
    "np.random.seed(1)\n",
    "\n",
    "# Create sample data.\n",
    "a = np.random.randn(40)\n",
    "b = 4*np.random.randn(50)\n",
    "\n",
    "# Use scipy.stats.ttest_ind.\n",
    "t, p = ttest_ind(a, b, equal_var=False)\n",
    "print(\"ttest_ind:            t = %g  p = %g\" % (t, p))\n",
    "\n",
    "# Compute the descriptive statistics of a and b.\n",
    "abar = a.mean()\n",
    "avar = a.var(ddof=1)\n",
    "na = a.size\n",
    "adof = na - 1\n",
    "\n",
    "bbar = b.mean()\n",
    "bvar = b.var(ddof=1)\n",
    "nb = b.size\n",
    "bdof = nb - 1\n",
    "\n",
    "# Use scipy.stats.ttest_ind_from_stats.\n",
    "t2, p2 = ttest_ind_from_stats(abar, np.sqrt(avar), na,\n",
    "                              bbar, np.sqrt(bvar), nb,\n",
    "                              equal_var=False)\n",
    "print(\"ttest_ind_from_stats: t = %g  p = %g\" % (t2, p2))\n",
    "\n",
    "# Use the formulas directly.\n",
    "tf = (abar - bbar) / np.sqrt(avar/na + bvar/nb)\n",
    "dof = (avar/na + bvar/nb)**2 / (avar**2/(na**2*adof) + bvar**2/(nb**2*bdof))\n",
    "pf = 2*stdtr(dof, -np.abs(tf))\n",
    "\n",
    "print(\"formula:              t = %g  p = %g\" % (tf, pf))\n",
    "\n",
    "a = torch.tensor(a)\n",
    "b = torch.tensor(b)\n",
    "tf = (a.mean() - b.mean()) / torch.sqrt(a.var()/a.size(0) + b.var()/b.size(0))\n",
    "print(\"formula:              t = %g\" % (tf))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class PositionwiseFeedForward(nn.Sequential):\n",
    "    def __init__(self, dim, dropout=0., act='reglu', mlp_ratio=1):\n",
    "        act_mult = 2 if act.lower() in [\"geglu\", \"reglu\"] else 1\n",
    "        super().__init__(nn.Linear(dim, dim * mlp_ratio * act_mult),\n",
    "                         get_act_fn(act),\n",
    "                         nn.Dropout(dropout),\n",
    "                         nn.Linear(dim * mlp_ratio, dim),\n",
    "                         nn.Dropout(dropout))\n",
    "\n",
    "class TokenLayer(Module):\n",
    "    def __init__(self, token=True): self.token = token\n",
    "    def forward(self, x): return x[..., 0] if self.token is not None else x.mean(-1)\n",
    "    def __repr__(self): return f\"{self.__class__.__name__}()\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = torch.randn(2,3,10)\n",
    "m = PositionwiseFeedForward(10, dropout=0., act='reglu', mlp_ratio=1)\n",
    "test_eq(m(t).shape, t.shape)\n",
    "m = PositionwiseFeedForward(10, dropout=0., act='smelu', mlp_ratio=1)\n",
    "test_eq(m(t).shape, t.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class ScaledDotProductAttention(Module):\n",
    "    r\"\"\"Scaled Dot-Product Attention module (Attention is all you need by Vaswani et al., 2017) with optional residual attention from previous layer\n",
    "    (Realformer: Transformer likes residual attention by He et al, 2020) and locality self sttention (Vision Transformer for Small-Size Datasets\n",
    "    by Lee et al, 2021)\"\"\"\n",
    "\n",
    "    def __init__(self, d_model, n_heads, attn_dropout=0., res_attention=False, lsa=False):\n",
    "        self.attn_dropout = nn.Dropout(attn_dropout)\n",
    "        self.res_attention = res_attention\n",
    "        head_dim = d_model // n_heads\n",
    "        self.scale = nn.Parameter(torch.tensor(head_dim ** -0.5), requires_grad=lsa)\n",
    "        self.lsa = lsa\n",
    "\n",
    "    def forward(self, q:Tensor, k:Tensor, v:Tensor, prev:Optional[Tensor]=None, key_padding_mask:Optional[Tensor]=None, attn_mask:Optional[Tensor]=None):\n",
    "        '''\n",
    "        Input shape:\n",
    "            q               : [bs x n_heads x max_q_len x d_k]\n",
    "            k               : [bs x n_heads x d_k x seq_len]\n",
    "            v               : [bs x n_heads x seq_len x d_v]\n",
    "            prev            : [bs x n_heads x q_len x seq_len]\n",
    "            key_padding_mask: [bs x seq_len]\n",
    "            attn_mask       : [1 x seq_len x seq_len]\n",
    "\n",
    "        Output shape:\n",
    "            output:  [bs x n_heads x q_len x d_v]\n",
    "            attn   : [bs x n_heads x q_len x seq_len]\n",
    "            scores : [bs x n_heads x q_len x seq_len]\n",
    "        '''\n",
    "\n",
    "        # Scaled MatMul (q, k) - similarity scores for all pairs of positions in an input sequence\n",
    "        attn_scores = torch.matmul(q, k) * self.scale      # attn_scores : [bs x n_heads x max_q_len x q_len]\n",
    "\n",
    "        # Add pre-softmax attention scores from the previous layer (optional)\n",
    "        if prev is not None: attn_scores = attn_scores + prev\n",
    "\n",
    "        # Attention mask (optional)\n",
    "        if attn_mask is not None:                                     # attn_mask with shape [q_len x seq_len] - only used when q_len == seq_len\n",
    "            if attn_mask.dtype == torch.bool:\n",
    "                attn_scores.masked_fill_(attn_mask, -np.inf)\n",
    "            else:\n",
    "                attn_scores += attn_mask\n",
    "\n",
    "        # Key padding mask (optional)\n",
    "        if key_padding_mask is not None:                              # mask with shape [bs x q_len] (only when max_w_len == q_len)\n",
    "            attn_scores.masked_fill_(key_padding_mask.unsqueeze(1).unsqueeze(2), -np.inf)\n",
    "\n",
    "        # normalize the attention weights\n",
    "        attn_weights = F.softmax(attn_scores, dim=-1)                 # attn_weights   : [bs x n_heads x max_q_len x q_len]\n",
    "        attn_weights = self.attn_dropout(attn_weights)\n",
    "\n",
    "        # compute the new values given the attention weights\n",
    "        output = torch.matmul(attn_weights, v)                        # output: [bs x n_heads x max_q_len x d_v]\n",
    "\n",
    "        if self.res_attention: return output, attn_weights, attn_scores\n",
    "        else: return output, attn_weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor(-2.3159e-10, grad_fn=<MeanBackward0>),\n",
       " tensor(0.9743, grad_fn=<StdBackward0>))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "B = 16\n",
    "C = 10\n",
    "M = 1500 # seq_len\n",
    "\n",
    "n_heads = 1\n",
    "D = 128 # model dimension\n",
    "N = 512 # max_seq_len - latent's index dimension\n",
    "d_k = D // n_heads\n",
    "\n",
    "xb = torch.randn(B, C, M)\n",
    "xb = (xb - xb.mean()) / xb.std()\n",
    "\n",
    "# Attention\n",
    "# input (Q)\n",
    "lin = nn.Linear(M, N, bias=False)\n",
    "Q = lin(xb).transpose(1,2)\n",
    "test_eq(Q.shape, (B, N, C))\n",
    "\n",
    "# q\n",
    "to_q = nn.Linear(C, D, bias=False)\n",
    "q = to_q(Q)\n",
    "q = nn.LayerNorm(D)(q)\n",
    "\n",
    "# k, v\n",
    "context = xb.transpose(1,2)\n",
    "to_kv = nn.Linear(C, D * 2, bias=False)\n",
    "k, v = to_kv(context).chunk(2, dim = -1)\n",
    "k = k.transpose(-1, -2)\n",
    "k = nn.LayerNorm(M)(k)\n",
    "v = nn.LayerNorm(D)(v)\n",
    "\n",
    "test_eq(q.shape, (B, N, D))\n",
    "test_eq(k.shape, (B, D, M))\n",
    "test_eq(v.shape, (B, M, D))\n",
    "\n",
    "output, attn, scores = ScaledDotProductAttention(D, n_heads, res_attention=True)(q.unsqueeze(1), k.unsqueeze(1), v.unsqueeze(1))\n",
    "test_eq(output.shape, (B, 1, N, D))\n",
    "test_eq(attn.shape, (B, 1, N, M))\n",
    "test_eq(scores.shape, (B, 1, N, M))\n",
    "scores.mean(), scores.std()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class MultiheadAttention(Module):\n",
    "    def __init__(self, d_model, n_heads, d_k=None, d_v=None, res_attention=False, attn_dropout=0., proj_dropout=0., qkv_bias=True, lsa=False):\n",
    "        \"\"\"Multi Head Attention Layer\n",
    "\n",
    "        Input shape:\n",
    "            Q:       [batch_size (bs) x max_q_len x d_model]\n",
    "            K, V:    [batch_size (bs) x q_len x d_model]\n",
    "            mask:    [q_len x q_len]\n",
    "        \"\"\"\n",
    "\n",
    "        d_k = ifnone(d_k, d_model // n_heads)\n",
    "        d_v = ifnone(d_v, d_model // n_heads)\n",
    "\n",
    "        self.n_heads, self.d_k, self.d_v = n_heads, d_k, d_v\n",
    "\n",
    "        self.W_Q = nn.Linear(d_model, d_k * n_heads, bias=qkv_bias)\n",
    "        self.W_K = nn.Linear(d_model, d_k * n_heads, bias=qkv_bias)\n",
    "        self.W_V = nn.Linear(d_model, d_v * n_heads, bias=qkv_bias)\n",
    "\n",
    "        # Scaled Dot-Product Attention (multiple heads)\n",
    "        self.res_attention = res_attention\n",
    "        self.sdp_attn = ScaledDotProductAttention(d_model, n_heads, attn_dropout=attn_dropout, res_attention=self.res_attention, lsa=lsa)\n",
    "\n",
    "        # Poject output\n",
    "        self.to_out = nn.Sequential(nn.Linear(n_heads * d_v, d_model), nn.Dropout(proj_dropout))\n",
    "\n",
    "\n",
    "    def forward(self, Q:Tensor, K:Optional[Tensor]=None, V:Optional[Tensor]=None, prev:Optional[Tensor]=None,\n",
    "                key_padding_mask:Optional[Tensor]=None, attn_mask:Optional[Tensor]=None):\n",
    "\n",
    "        bs = Q.size(0)\n",
    "        if K is None: K = Q\n",
    "        if V is None: V = Q\n",
    "\n",
    "        # Linear (+ split in multiple heads)\n",
    "        q_s = self.W_Q(Q).reshape(bs, -1, self.n_heads, self.d_k).transpose(1,2)       # q_s    : [bs x n_heads x max_q_len x d_k]\n",
    "        k_s = self.W_K(K).reshape(bs, -1, self.n_heads, self.d_k).permute(0,2,3,1)     # k_s    : [bs x n_heads x d_k x q_len] - transpose(1,2) + transpose(2,3)\n",
    "        v_s = self.W_V(V).reshape(bs, -1, self.n_heads, self.d_v).transpose(1,2)       # v_s    : [bs x n_heads x q_len x d_v]\n",
    "\n",
    "        # Apply Scaled Dot-Product Attention (multiple heads)\n",
    "        if self.res_attention:\n",
    "            output, attn_weights, attn_scores = self.sdp_attn(q_s, k_s, v_s, prev=prev, key_padding_mask=key_padding_mask, attn_mask=attn_mask)\n",
    "        else:\n",
    "            output, attn_weights = self.sdp_attn(q_s, k_s, v_s, key_padding_mask=key_padding_mask, attn_mask=attn_mask)\n",
    "        # output: [bs x n_heads x q_len x d_v], attn: [bs x n_heads x q_len x q_len], scores: [bs x n_heads x max_q_len x q_len]\n",
    "\n",
    "        # back to the original inputs dimensions\n",
    "        output = output.transpose(1, 2).contiguous().reshape(bs, -1, self.n_heads * self.d_v) # output: [bs x q_len x n_heads * d_v]\n",
    "        output = self.to_out(output)\n",
    "\n",
    "        if self.res_attention: return output, attn_weights, attn_scores\n",
    "        else: return output, attn_weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "attn_mask torch.Size([50, 50]) key_padding_mask torch.Size([16, 50])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(torch.Size([16, 3, 50, 6]), torch.Size([16, 3, 50, 50]))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "q = torch.rand([16, 3, 50, 8])\n",
    "k = torch.rand([16, 3, 50, 8]).transpose(-1, -2)\n",
    "v = torch.rand([16, 3, 50, 6])\n",
    "attn_mask = torch.triu(torch.ones(50, 50)) # shape: q_len x q_len\n",
    "key_padding_mask = torch.zeros(16, 50)\n",
    "key_padding_mask[[1, 3, 6, 15], -10:] = 1\n",
    "key_padding_mask = key_padding_mask.bool()\n",
    "print('attn_mask', attn_mask.shape, 'key_padding_mask', key_padding_mask.shape)\n",
    "output, attn = ScaledDotProductAttention(24, 3, attn_dropout=.1)(q, k, v, attn_mask=attn_mask, key_padding_mask=key_padding_mask)\n",
    "output.shape, attn.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([16, 50, 128]), torch.Size([16, 3, 50, 50]))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "t = torch.rand(16, 50, 128)\n",
    "output, attn = MultiheadAttention(d_model=128, n_heads=3, d_k=8, d_v=6)(t, t, t, key_padding_mask=key_padding_mask, attn_mask=attn_mask)\n",
    "output.shape, attn.shape"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Test multi-head attention with self-locality attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([16, 50, 128]), torch.Size([16, 8, 50, 50]))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# lsa (locality self-sttention)\n",
    "t = torch.rand(16, 50, 128)\n",
    "attn_mask = torch.eye(50).reshape(1, 1, 50, 50).bool()\n",
    "output, attn = MultiheadAttention(d_model=128, n_heads=8, lsa=True)(t, t, t, key_padding_mask=key_padding_mask, attn_mask=attn_mask)\n",
    "output.shape, attn.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = torch.rand(16, 50, 128)\n",
    "att_mask = (torch.rand((50, 50)) > .85).float()\n",
    "att_mask[att_mask == 1] = -np.inf\n",
    "\n",
    "mha = MultiheadAttention(d_model=128, n_heads=3, d_k=8, d_v=6)\n",
    "output, attn = mha(t, t, t, attn_mask=att_mask)\n",
    "test_eq(torch.isnan(output).sum().item(), 0)\n",
    "test_eq(torch.isnan(attn).sum().item(), 0)\n",
    "loss = output[:2, :].sum()\n",
    "test_eq(torch.isnan(loss).sum().item(), 0)\n",
    "loss.backward()\n",
    "for n, p in mha.named_parameters():\n",
    "    if p.grad is not None:\n",
    "        test_eq(torch.isnan(p.grad).sum().item(), 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = torch.rand(16, 50, 128)\n",
    "attn_mask = (torch.rand((50, 50)) > .85)\n",
    "\n",
    "# True values will be masked\n",
    "mha = MultiheadAttention(d_model=128, n_heads=3, d_k=8, d_v=6)\n",
    "output, attn = mha(t, t, t, attn_mask=att_mask)\n",
    "test_eq(torch.isnan(output).sum().item(), 0)\n",
    "test_eq(torch.isnan(attn).sum().item(), 0)\n",
    "loss = output[:2, :].sum()\n",
    "test_eq(torch.isnan(loss).sum().item(), 0)\n",
    "loss.backward()\n",
    "for n, p in mha.named_parameters():\n",
    "    if p.grad is not None:\n",
    "        test_eq(torch.isnan(p.grad).sum().item(), 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class MultiConv1d(Module):\n",
    "    \"\"\"Module that applies multiple convolutions with different kernel sizes\"\"\"\n",
    "\n",
    "    def __init__(self, ni, nf=None, kss=[1,3,5,7], keep_original=False, separable=False, dim=1, **kwargs):\n",
    "        kss = listify(kss)\n",
    "        n_layers = len(kss)\n",
    "        if ni == nf: keep_original = False\n",
    "        if nf is None: nf = ni * (keep_original + n_layers)\n",
    "        nfs = [(nf - ni*keep_original) // n_layers] * n_layers\n",
    "        while np.sum(nfs) + ni * keep_original < nf:\n",
    "            for i in range(len(nfs)):\n",
    "                nfs[i] += 1\n",
    "                if np.sum(nfs) + ni * keep_original == nf: break\n",
    "\n",
    "        _conv = SeparableConv1d if separable else Conv1d\n",
    "        self.layers = nn.ModuleList()\n",
    "        for nfi,ksi in zip(nfs, kss):\n",
    "            self.layers.append(_conv(ni, nfi, ksi, **kwargs))\n",
    "        self.keep_original, self.dim = keep_original, dim\n",
    "\n",
    "    def forward(self, x):\n",
    "        output = [x] if self.keep_original else []\n",
    "        for l in self.layers:\n",
    "            output.append(l(x))\n",
    "        x = torch.cat(output, dim=self.dim)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = torch.rand(16, 6, 37)\n",
    "test_eq(MultiConv1d(6, None, kss=[1,3,5], keep_original=True)(t).shape, [16, 24, 37])\n",
    "test_eq(MultiConv1d(6, 36, kss=[1,3,5], keep_original=False)(t).shape, [16, 36, 37])\n",
    "test_eq(MultiConv1d(6, None, kss=[1,3,5], keep_original=True, dim=-1)(t).shape, [16, 6, 37*4])\n",
    "test_eq(MultiConv1d(6, 60, kss=[1,3,5], keep_original=True)(t).shape, [16, 60, 37])\n",
    "test_eq(MultiConv1d(6, 60, kss=[1,3,5], separable=True)(t).shape, [16, 60, 37])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class LSTMOutput(Module):\n",
    "    def forward(self, x): return x[0]\n",
    "    def __repr__(self): return f'{self.__class__.__name__}()'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = ([1], [2], [3])\n",
    "test_eq(LSTMOutput()(t), [1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def emb_sz_rule(n_cat):\n",
    "    \"Rule of thumb to pick embedding size corresponding to `n_cat` (original from fastai)\"\n",
    "    return min(600, round(1.6 * n_cat**0.56))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(emb_sz_rule(7), 5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class TSEmbedding(nn.Embedding):\n",
    "    \"Embedding layer with truncated normal initialization adapted from fastai\"\n",
    "    def __init__(self, ni, nf, std=0.01, padding_idx=None):\n",
    "        super().__init__(ni, nf)\n",
    "        trunc_normal_(self.weight.data, std=std)\n",
    "        if padding_idx is not None:\n",
    "            nn.init.zeros_(self.weight.data[padding_idx])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class MultiEmbedding(Module):\n",
    "    def __init__(self, c_in, n_cat_embeds, cat_embed_dims=None, cat_pos=None, std=0.01, cat_padding_idxs=None):\n",
    "        cat_n_embeds = listify(n_cat_embeds)\n",
    "        if cat_padding_idxs is None: cat_padding_idxs = [None]\n",
    "        else: cat_padding_idxs = listify(cat_padding_idxs)\n",
    "        if len(cat_padding_idxs) == 1 and len(cat_padding_idxs) < len(cat_n_embeds):\n",
    "            cat_padding_idxs = cat_padding_idxs * len(cat_n_embeds)\n",
    "        assert len(cat_n_embeds) == len(cat_padding_idxs)\n",
    "        if cat_embed_dims is None:\n",
    "            cat_embed_dims = [emb_sz_rule(s) for s in cat_n_embeds]\n",
    "        else:\n",
    "            cat_embed_dims = listify(cat_embed_dims)\n",
    "            if len(cat_embed_dims) == 1: cat_embed_dims = cat_embed_dims * len(cat_n_embeds)\n",
    "            assert len(cat_embed_dims) == len(cat_n_embeds)\n",
    "        if cat_pos:\n",
    "            cat_pos = torch.as_tensor(listify(cat_pos))\n",
    "        else:\n",
    "            cat_pos = torch.arange(len(cat_n_embeds))\n",
    "        self.register_buffer(\"cat_pos\", cat_pos)\n",
    "        cont_pos = torch.tensor([p for p in torch.arange(c_in) if p not in self.cat_pos])\n",
    "        self.register_buffer(\"cont_pos\", cont_pos)\n",
    "        self.cat_embed = nn.ModuleList([TSEmbedding(n,d,std=std, padding_idx=p) for n,d,p in zip(cat_n_embeds, cat_embed_dims, cat_padding_idxs)])\n",
    "\n",
    "    def forward(self, x):\n",
    "        if isinstance(x, tuple): x_cat, x_cont, *_ = x\n",
    "        else: x_cat, x_cont = x[:, self.cat_pos], x[:, self.cont_pos]\n",
    "        x_cat = torch.cat([e(torch.round(x_cat[:,i]).long()).transpose(1,2) for i,e in enumerate(self.cat_embed)],1)\n",
    "        return torch.cat([x_cat, x_cont], 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[3, 4] [3, 3] torch.Size([4, 3, 10]) torch.Size([4, 7, 10])\n"
     ]
    }
   ],
   "source": [
    "a = alphabet[np.random.randint(0,3,40)]\n",
    "b = ALPHABET[np.random.randint(6,10,40)]\n",
    "c = np.random.rand(40).reshape(4,1,10)\n",
    "map_a = {k:v for v,k in enumerate(np.unique(a))}\n",
    "map_b = {k:v for v,k in enumerate(np.unique(b))}\n",
    "n_embeds = [len(m.keys()) for m in [map_a, map_b]]\n",
    "szs = [emb_sz_rule(n) for n in n_embeds]\n",
    "a = np.asarray(a.map(map_a)).reshape(4,1,10)\n",
    "b = np.asarray(b.map(map_b)).reshape(4,1,10)\n",
    "inp = torch.from_numpy(np.concatenate((c,a,b), 1)).float()\n",
    "memb = MultiEmbedding(3, n_embeds, cat_pos=[1,2])\n",
    "# registered buffers are part of the state_dict() but not module.parameters()\n",
    "assert all([(k in memb.state_dict().keys()) for k in ['cat_pos', 'cont_pos']])\n",
    "embeddings = memb(inp)\n",
    "print(n_embeds, szs, inp.shape, embeddings.shape)\n",
    "test_eq(embeddings.shape, (inp.shape[0],sum(szs)+1,inp.shape[-1]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "me = MultiEmbedding(3, 4, cat_pos=2)\n",
    "test_eq(me.cat_embed[0].weight.shape, (4,3))\n",
    "test_eq(me.cat_pos.cpu().item(), 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/javascript": "IPython.notebook.save_checkpoint();",
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/Users/nacho/notebooks/tsai/nbs/029_models.layers.ipynb saved at 2025-01-22 18:40:22\n",
      "Correct notebook to script conversion! 😃\n",
      "Wednesday 22/01/25 18:40:25 CET\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "                <audio  controls=\"controls\" autoplay=\"autoplay\">\n",
       "                    <source src=\"data:audio/wav;base64,UklGRvQHAABXQVZFZm10IBAAAAABAAEAECcAACBOAAACABAAZGF0YdAHAAAAAPF/iPh/gOoOon6w6ayCoR2ZeyfbjobxK+F2Hs0XjKc5i3DGvzaTlEaraE+zz5uLUl9f46fHpWJdxVSrnfmw8mYEScqUP70cb0Q8X41uysJ1si6Eh1jYzXp9IE2DzOYsftYRyoCY9dJ/8QICgIcEun8D9PmAaBPlfT7lq4MFIlh61tYPiCswIHX+yBaOqT1QbuW7qpVQSv9lu6+xnvRVSlyopAypbGBTUdSalrSTaUBFYpInwUpxOzhti5TOdndyKhCGrdwAfBUcXIJB69p+Vw1egB76+n9q/h6ADglbf4LvnIHfF/981ODThF4m8HiS0riJVjQ6c+/EOZCYQfJrGrhBmPVNMmNArLKhQlkXWYqhbaxXY8ZNHphLuBJsZUEckCTFVHMgNKGJytIDeSUmw4QN4Qx9pReTgb3vYX/TCBuApf75f+P5Y4CRDdN+B+tngk8c8nt03CKGqipgd13OhotwOC5x9MCAknFFcmlmtPmagFFFYOCo0qRzXMhVi57pryNmIEqJlRi8bm52PfuNM8k4dfQv+4cO12l6zCGdg3jl730uE/KAPvS+f0wEAoAsA89/XfXQgBESIn6S5luDtiC8eh/YmIfpLqt1OMp5jXg8/24MveqUNUnPZsqw0Z3yVDldnaUOqIZfXlKrm36zzWhjRhaT+r+ncHI5/otUzfd2uSt7hl/bqXtoHaCC6+mqfrAOeoDD+PJ/xf8RgLMHfH/b8GeBihZIfSXidoQSJWB52NM1iRkzz3MkxpKPbUCrbDu5d5fgTAxkSK3JoEhYD1p2omere2LZTuqYLbdWa49Cx5Dww7tyXDUnioXRkHhwJyKFvd/AfPoYy4Fl7j1/LQorgEr9/X89+0qAOAwAf13sJoL8Gkd8wt25hWIp3Heez/eKODfPcSPCzpFNRDVqf7UlmnNQKGHgqd+jgVvJVm2f265QZTpLS5byur1tpT6ajvrHq3Q2MXWIxtUCehoj8YMk5LB9hRQegeTypn+nBQWA0QHgf7f2q4C5EFt+5ucOg2YfHXtq2SSHpS0ydnTL4IxFO6pvNb4ulBdInWfcsfSc7VMmXpSmE6eeXmZThJxpsgRohEfOk86+AHCoOpOMFsx1dv8s6oYT2k17uR7ngpXod34IEJqAaPfnfyABCIBZBpl/NPI2gTQVjX134x2ExSPMeR7VtYjZMWJ0W8ftjkA/YW1durCWykvjZFKu4p9LVwVbZKNkqpxh6U+6mRC2mGq2Q3SRvsIgcpc2sIpD0Bp4uiiFhW3ecXxOGgaCDe0Vf4cLPoDv+/5/mfw1gN4KKX+17emBqBmYfBHfVYUZKFR44NBtiv41bHJUwx+RJkP1apu2VJlkTwli4qrwoo1ax1dToNCtemRSTBGXz7kJbdM/PY/Dxht0dTLziH7Ul3loJEiE0uJsfdsVTYGL8Yt/AgcMgHYA7X8S+IqAYA+QfjzpxIIVHnp7tdqzhmAstXaxzEqMETpScGC/dJP3Rmdo8LIZnOVSEF+Opxumsl1sVF+dVrE5Z6NIiZSkvVdv2zsqjdnK8HVDLlyHyNjuegogM4NA5z9+YRG9gA722H97AgOA/gSyf43zCIHdE899yuTIg3ciNXpm1jmImTDwdJPITI4RPhRugbvslbFKt2Vfr/6eTFb4W1WkY6m6YPdQjJr2tNZp3EQlko7BgXHRNz2LAc+gdwMq7IUf3R58ohtFgrbr6n7hDFWAlPr8f/T9I4CECU9/De+vgVQY5nxh4POEzybJeCTS5YnCNAZzhsRzkP1Bsmu4t4aYU07nYuerA6KWWcJYO6HHrKJjaE3Zl624UWz/QOOPjcWHc7QzdIk40yl5tCWjhIDhJX0xF4CBMvBsf10IF4Ac//Z/bPlsgAcOwn6S6n6CwxzUewLcRoYaKzV38M23i9o493CNwL6S1UUuaQe0QpvbUfdfiqglpcRccFU+nkWwambASUiVfLyqbg49xY2eyWh1hy/Sh37XjHpaIYKD7OUEfrgS5IC09MV/1gMBgKMDyH/n9N6AhhINfh7mdoMoIZt6r9fAh1cvfHXNya6N4DzDbqi8K5WWSYlmbbAdnkpV6FxJpWSo1V8DUmGb3rMRaQBG2JJgwN9wCDnNi8HNI3dKK1aG0dvHe/UciIJf6rt+Og5wgDn59X9P/xWAKQhxf2XweYH+FjB9suGVhIMlOnlo02GJhTOdc7vFyo/TQGxs2Li7lz9NwmPurBihnVi7WSWiwKvGYntOpJiOt5drKUKMkFnE8HLxNPmJ9NG4eP8mAYUv4Np8hhi3gdruSX+3CSWAwP38f8f6UoCuDPF+6Os8gnAbKnxQ3d2F0imydzDPKIuiN5lxu8EKkrFE82kftW2az1DbYImpMqTUW3FWIJ83r5hl2koJlla7+m0+PmSOZcjcdMgwS4g11iZ6qCLUg5jkxn0QFA6BWvOvfzEFBIBHAtp/Qfa3gC4RSH5y5yeD2B/8evnYS4cULgR2CMsUja47cG/QvW6UeEhXZ3+xP51GVNVdP6Zpp+1eDFM5nMeySWghR4+TNL85cD46YIyCzKJ2kCzEhoTabXtGHs+CCemJfpMPjoDe9+t/qQALgM8Gj3++8UaBqRV2fQTjO4Q3JKd5r9TgiEYyMHTxxiWPpz8jbfq585YpTJpk960xoKFXsVoTo7yq6GGMTw==\" type=\"audio/wav\" />\n",
       "                    Your browser does not support the audio element.\n",
       "                </audio>\n",
       "              "
      ],
      "text/plain": [
       "<IPython.lib.display.Audio object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#|eval: false\n",
    "#|hide\n",
    "from tsai.export import get_nb_name; nb_name = get_nb_name(locals())\n",
    "from tsai.imports import create_scripts; create_scripts(nb_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
